adding static_keynames support

[mishmish66]

Added support for static_keynames which lets you specify key names to ignore for jax transformations #64 , all tests seem to be passing I've attached the test log (bash test.sh > test_out.txt) test_out.txt

I'd love to have a cleaner way to get the static_keynames in the flatten func if anyone has ideas, this way there's no indication to users that static_keynames is something placed in there by chex, but calling it _static_keynames set off the tests for private access.

To demonstrate a use case with this I can dataclassify this:

@chex.dataclass(frozen=True)
class Foo:
    name: str
    data: jnp.ndarray

def noise(key, foo):
    rng, key = jax.random.split(key)
    noise = jax.random.normal(rng, foo.data.shape)
    return foo.data + noise

but when I try to JIT the noise function it will crash because the name in the dataclass is a str which is not a valid jax type.

If I do this instead:

@chex.dataclass(frozen=True, static_keynames=["name"])
class Foo:
    name: str
    data: jnp.ndarray

def noise(key, foo):
    rng, key = jax.random.split(key)
    noise = jax.random.normal(rng, foo.data.shape)
    return foo.data + noise

then it won't crash since it will make the name field static now I can jit it just fine (but if I try to use a static field in a jax transformed function in a non-static way then it will crash)

[google-cla[bot]]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[mishmish66]

Is there anything else I should do to get this change in?