Extend use_parameter_layout_on_device option to ExecuteReplicated.

`ExecuteReplicated` ignored the `use_parameter_layout_on_device` option of the
`HloRunnerPjRt`. If enabled, this flag passes the parameter layout for use with
the on-device buffer. This is something that `HloRunner` does as well, so this
functionality is just replicating the existing behavior for use with PjRt.

PiperOrigin-RevId: 706080433

xla/service/hlo_runner_pjrt.cc

@@ -480,8 +480,12 @@ absl::StatusOr<std::vector<Literal>> HloRunnerPjRt::ExecuteReplicatedImpl(
       const Literal* const argument = argument_provider(i, arg_index);
       TF_RET_CHECK(argument != nullptr);
 
-      TF_ASSIGN_OR_RETURN(auto assignment, pjrt_client_->BufferFromHostLiteral(
-                                               *argument, device_ptr));
+      TF_ASSIGN_OR_RETURN(
+          std::unique_ptr<PjRtBuffer> assignment,
+          use_parameter_layout_on_device_
+              ? pjrt_client_->BufferFromHostLiteral(*argument, device_ptr,
+                                                    &argument->shape().layout())
+              : pjrt_client_->BufferFromHostLiteral(*argument, device_ptr));
       replica_buffers.push_back(std::move(assignment));
     }