Mahmoud Soliman

Have the same issue after manually changing ```python data= layers.Input(name = 'data', shape = (112, 112, 3,) ) # The following line throws could not find minusscalar0_second minusscalar0 = my_sub()(data,...

After reading this [This issue](https://github.com/keras-team/keras/issues/11811) I figured out that everything needs to be a keras Layer so I changed those lines to : ```python data = layers.Input(name='data', shape=(112, 112, 3,))...

I think this would be much better using the GitHub projects / boards feature :)

@avital @levskaya anything needed from my side to merge this?

@sudhakarsingh27 @nvcastet @nouiz for viz

NVIDIA official images are coming soon 😄 @pwais @lukaemon

The NVIDIA JAX early access (EA) container is now released, [register ](https://developer.nvidia.com/jax-container-early-access)to get it :)

CC: @nouiz

> > General comment: should we perhaps choose a simpler op/ops for this? RMSNorm needs a lot of CUDA code which is not directly relevant here. > > People don't...

@aterenin, even though it seems like an XLA bug, it would be helpful to mention which hardware you are using (compilers have different backends, so code paths are different).