Update JAX tests for separate input_striding and input_tiling on transposes.

This updates JAX after an PJRT API change.

PiperOrigin-RevId: 846260872

Author: hawkinsp

jaxlib/callback.cc

@@ -104,7 +104,7 @@ absl::Status CpuCallback::PrepareAndCall(void** result, void** arg_ptrs) {
           xla::primitive_util::ByteWidth(results_[i].type);
       options.dims = dims;
       options.permutation = results_[i].reversed_layout;
-      options.input_layout = xla::TransposePlan::Striding{strides};
+      options.input_striding = xla::TransposePlan::Striding{strides};
       absl::StatusOr<std::shared_ptr<xla::TransposePlan>> plan =
           transpose_cache_.GetOrCreate(options);
       if (!plan.ok()) {

jaxlib/gpu/py_client_gpu.cc

@@ -217,7 +217,7 @@ xla::ffi::Error XlaFfiPythonGpuCallback(gpuStream_t stream,
       absl::c_reverse_copy(expected_shape.layout().minor_to_major(),
                            reversed_layout.begin());
       options.permutation = reversed_layout;
-      options.input_layout = xla::TransposePlan::Striding{strides};
+      options.input_striding = xla::TransposePlan::Striding{strides};
       auto maybe_plan = transpose_cache->cache.GetOrCreate(options);
       if (!maybe_plan.ok()) {
         return xla::ffi::Error::Internal(maybe_plan.status().ToString());

jaxlib/py_client_cpu.cc

@@ -173,7 +173,7 @@ ffi::Error XlaFfiPythonCpuCallback(xla::FfiLoadedHostCallbacks* callbacks,
       absl::c_reverse_copy(expected_shape.layout().minor_to_major(),
                            reversed_layout.begin());
       options.permutation = reversed_layout;
-      options.input_layout = xla::TransposePlan::Striding{strides};
+      options.input_striding = xla::TransposePlan::Striding{strides};
       auto maybe_plan = transpose_cache->cache.GetOrCreate(options);
       if (!maybe_plan.ok()) {
         return ffi::Error::Internal(maybe_plan.status().ToString());