dtensor.run_on() only accept Mesh
dtensor.run_on() only accept Mesh

Removal of layout_or: the result does not ensure the tensors follow layout,
passing in a Layout is confusing.

Existing users of layout_or path can migrate to to:

  dtensor.run_on(layout.mesh)

  Also, to use the context conditionally:

  with dtensor.run_on(mesh) if mesh else contextlib.nullcontext():

PiperOrigin-RevId: 443454416

keras-team/keras,keras-team/keras