Function fn

axlearn/common/input_base_test.py:126–129

(input_batch)

    @partial(pjit, in_shardings=None)
    def fn(input_batch):
        output_batch = ds.dispatch_global_batch(input_batch)
        jax.tree.map(check_sharding, tree_paths(output_batch), output_batch)
        return output_batch

    fn.lower(input_batch).compile()