Automated Code Change

PiperOrigin-RevId: 872765836

Author: Google-ML-Automation

docs/pallas/CHANGELOG.md

@@ -25,6 +25,8 @@ Remember to align the itemized text with the first line of an item within a list
     For example, to force the use of the Triton backend you have to now write
     `compiler_params=pltriton.CompilerParams()`, where `pltriton` refers to
     {mod}`jax.experimental.pallas.triton`.
+  * Renamed {class}`jax.experimental.pallas.tpu.KernelType` to `CoreType`. The
+    old name is deprecated and will be removed in a future release.
 
 ## Released with JAX 0.9.0
 

jax/_src/pallas/mosaic/core.py

@@ -43,7 +43,7 @@
 _out_shape_to_aval_mapping = pallas_core._out_shape_to_aval_mapping
 
 
-class KernelType(enum.Enum):
+class CoreType(enum.Enum):
   TC = 0
   SC_SCALAR_SUBCORE = 1
   SC_VECTOR_SUBCORE = 2
@@ -114,7 +114,7 @@ class CompilerParams(pallas_core.CompilerParams):
   flags: dict[str, Any] | None = None
   internal_scratch_in_bytes: int | None = None
   serialization_format: int = 1
-  kernel_type: KernelType = KernelType.TC
+  kernel_type: CoreType = CoreType.TC
   disable_bounds_checks: bool = False
   skip_device_barrier: bool = False
   allow_collective_id_without_custom_barrier: bool = False
@@ -131,7 +131,7 @@ def __init__(
       flags: Mapping[str, Any] | None = None,
       internal_scratch_in_bytes: int | None = None,
       serialization_format: int = 1,
-      kernel_type: KernelType = KernelType.TC,
+      kernel_type: CoreType = CoreType.TC,
       disable_bounds_checks: bool = False,
       skip_device_barrier: bool = False,
       allow_collective_id_without_custom_barrier: bool = False,
@@ -190,7 +190,7 @@ def __str__(self) -> str:
   def from_type(self, ty):
     return pallas_core.MemoryRef(ty, memory_space=self)
 
-  def __call__(self, shape: Sequence[int], dtype: jnp.dtype):
+  def __call__(self, shape: Sequence[int], dtype: jnp.dtype[Any]):
     # A convenience function for constructing MemoryRef types of ShapedArrays.
     return self.from_type(jax_core.ShapedArray(tuple(shape), dtype))
 
@@ -283,8 +283,8 @@ def __hash__(self) -> int:
     )
 
   @property
-  def kernel_type(self) -> KernelType:
-    return KernelType.TC
+  def kernel_type(self) -> CoreType:
+    return CoreType.TC
 
   @property
   def default_memory_space(self) -> pallas_core.MemorySpace: