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

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

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

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

import argparse

import jax
from jax import lax, numpy as jnp, random
from matplotlib import pyplot as plt, rcParams


def newton_raphson(f, x, iters):
    """Use the Newton-Raphson method to find a root of the given function."""

carlosgmartin

import jax
from jax import numpy as jnp
from jax.flatten_util import ravel_pytree


def centered_finite_differences(f, x, eps):
    x, unravel = ravel_pytree(x)

    def helper(index):
        xp = x.at[index].add(eps)

carlosgmartin

from functools import partial

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


def pseudo_gradient(
    f,
    x,

carlosgmartin

import numpy as np
import gurobipy
import cvxopt
from multiprocessing import Pool
from timeit import default_timer as timer

gurobipy.setParam('OutputFlag', 0)
def max_min_gurobi(game):
    model = gurobipy.Model()
    v = model.addMVar(1)

carlosgmartin

import numpy as np
import sympy.ntheory
import time

# https://mathoverflow.net/a/227408/74578
def mobius_transform(sequence):
    sequence = sequence.copy()
    for i in range(1, len(sequence)//2+1):
        sequence[i+i-1::i] -= sequence[i-1]
    return sequence

carlosgmartin

# https://www.sciencedirect.com/science/article/pii/0304397588900461

def stackermann(i, n):
	stack = []
	stack.append(i)
	stack.append(n)
	while len(stack) > 1:
		n = stack.pop()
		i = stack.pop()
		if i == 0: