[Pallas-Triton] Fix squeeze lowering required sharding argument

[shangz-ai]

Error comes from this PR: #25103 if there is no sharding argument passed to squeeze.

I observed my pallas kernel (using squeeze) failure due to the above PR, and the kernel was fine before that PR was merged.

[shangz-ai]

@superbobry Here is a very naive reproducer

import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl
from functools import partial

def squeeze_kernel(k_ref, o_ref):
    # Read inputs
    k = k_ref[...]

    # Squeeze and store
    o_ref[...] = jnp.squeeze(k)

@partial(jax.jit, static_argnames='out_type')
def squeeze_k(k: jax.Array, out_type=jnp.float32) -> jax.Array:
    return pl.pallas_call(
        squeeze_kernel,
        out_shape=jax.ShapeDtypeStruct(jnp.squeeze(k).shape, out_type)
    )(k)

if __name__ == "__main__":
    k = jnp.ones((4, 2, 1, 1, 2, 1, 4), dtype=jnp.float32)

    # Run the kernel and print result
    print("Shape before squeezing: \n", k.shape)
    res = squeeze_k(k, out_type=jnp.float32)
    print("Shape after squeezing: \n", res.shape)

The main branch jax will report error as:

msg=_reshape_lowering_rule() missing 1 required keyword-only argument: 'sharding'

but with the PR fix, it can run correctly as:

Shape before squeezing: 
 (4, 2, 1, 1, 2, 1, 4)
Shape after squeezing: 
 (4, 2, 2, 4)

Please let me know if that makes sense to you or need any other information.
Thanks!

[superbobry]

Can you squash the commits please?

[shangz-ai]

Can you squash the commits please?

Done