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]

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?

[wtedw]

Any updates on this static_keynames feature, @hbq1 or whoever maintains this repo?
This would be very helpful for my projects.

[stompchicken]

Hi there @mishmish66, @wtedw. Apologies for ignoring this PR. @hbq1 and I have been quite negligent in maintaining Chex these days.

Thanks a lot for taking a pass at this. It is a commonly requested feature and I think it would be broadly useful. I would be happy for this to be merged, but it does need to have decent unit test coverage before that can happen.

[mishmish66]

I'll take a look at this later today! Thanks for getting around to it. I'll need to remember what is going on here (forgot about this PR) and I can run the tests.

[arik-shurygin]

Hey all, wanted to check in here. Did you get around to writing tests @mishmish66 ? Thanks for the work on this regardless, I know its an old PR at this point.

[mishmish66]

After recovering from the initial shock of reading my old code I took a look again and added some tests + cleaned it up a bit.

All the existing tests still passed so it seems not to break anything. I rebased and force pushed it to avoid a nasty merge in the history from an ancient branch.

Hope this works for everyone!

[mishmish66]

@stompchicken Do you think that the coverage is enough to make it in?