carlosgmartin

import argparse

import ale_py
import cv2
import gymnasium as gym
import minari
import minigrid
import numpy as np
import sb3_contrib
import shortuuid

carlosgmartin

from timeit import default_timer

import jax
import optax
import tqdm
from jax import lax, numpy as jnp, random


@jax.jit
def scenic_hungarian_algorithm(cost):

carlosgmartin

"""
Improving the lowering and compilation of unrolled lax.scan loops
https://github.com/jax-ml/jax/discussions/25336
"""

import argparse
import functools

import jax
from jax import lax, numpy as jnp, random

carlosgmartin

"""
https://github.com/jax-ml/jax/issues/26566
"""

import itertools
import operator

import jax
from jax import numpy as jnp
from tqdm import tqdm