Very slow compilation of game tree policy value gradient

[carlosgmartin]

Consider the following program (which uses OpenSpiel):

from time import time
from pyspiel import load_game
from jax import numpy as jnp, random, jit, grad, nn

def get_value(state, policy):
    if state.is_terminal():
        return jnp.asarray(state.returns())
    elif state.is_chance_node():
        return sum(
            prob * get_value(state.child(action), policy)
            for action, prob in state.chance_outcomes()
        )
    else:
        return sum(
            prob * get_value(state.child(action), policy)
            for action, prob in zip(state.legal_actions(), policy(state))
        )

def traverse(state):
    yield state
    for action in state.legal_actions():
        yield from traverse(state.child(action))

def init_params(key, root):
    info_actions = {
        state.information_state_string(): len(state.legal_actions())
        for state in traverse(root)
        if state.is_player_node()
    }
    return {
        info: random.normal(key, [actions])
        for (info, actions), subkey in
        zip(info_actions.items(), random.split(key, len(info_actions)))
    }

def get_payoffs(params, root):
    probs = {
        info: nn.softmax(logits)
        for info, logits in params.items()
    }
    def policy(state):
        return probs[state.information_state_string()]
    return get_value(root, policy)

def main():
    root = load_game('leduc_poker(suit_isomorphism=True)').new_initial_state()
    def f(params):
        return get_payoffs(params, root)[0]
    params = init_params(random.PRNGKey(0), root)

    start = time()
    grad(f)(params)
    print(time() - start)

    start = time()
    grad(jit(f))(params)
    print(time() - start)

    grad_f = jit(grad(f))
    for _ in range(3):
        start = time()
        grad_f(params)
        print(time() - start)

if __name__ == '__main__':
    main()

It outputs

5.8327319622039795
2022-10-10 18:39:01.455396: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65] 
********************************
[Compiling module jit_backward_pass] Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************
2022-10-10 18:39:03.185329: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 2m1.731974s

********************************
[Compiling module jit_backward_pass] Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************
175.74559116363525
2022-10-10 18:42:48.564768: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 3m14.714633s

********************************
[Compiling module jit_f] Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************
225.3683431148529
0.003172159194946289
0.0023038387298583984

My specs are the following:

Python 3.10.7
jax 0.3.21
jaxlib 0.3.20
macOS 11.7
2.5 GHz Quad-Core Intel Core i7
16 GB 1600 MHz DDR3

Is there any way to speed up the very slow compilation? More generally, is there any way to adjust the optimizations that are attempted by the compiler, in order to speed it up? (There is an old discussion related to this at #3478.)