jax

d24ffbd0 - [Pallas] Add .like constructor method for MemoryRefs in Pallas TPU and GPU

[Pallas] Add .like constructor method for MemoryRefs in Pallas TPU and GPU This method allows constructing a MemoryRef from a shape-dtype-like object. Example: pltpu.VMEM.like(x) and plgpu.GMEM.like(x) PiperOrigin-RevId: 907303922