Roger-luo

#################
# Utils
#################

"""
    log2i(x)

logrithm for integer pow of 2
"""
function log2i(x::T)::T where T

arthurmensch

"""
Custom twice differentiable functions in Pytorch.

Author: Arthur Mensch
"""


import torch

import torch.nn as nn

kylemcdonald

kylemcdonald

mattjj

from jax.interpreters import ad
from jax.interpreters import partial_eval as pe
from jax import custom_transforms
from jax import core
from jax import grad

@custom_transforms
def f(x, y):
  return x**2 + 3 * y

GiggleLiu

using TupleTools
using Base.Cartesian
using CuArrays, CUDAnative

"""
A naive implementation of `einsum!`
    * `ixs`: input tensor indices,
    * `xs`: input tensors,
    * `iy`: output tensor indices,
    * `y`: accumulated tensor, notice it is initialized to 0 as output!

caryan

shoyer

shoyer

Roger-luo

macro _threads(ex)
    return quote
        if (Threads.nthreads() > 1) && (length(st) > 4096)
            $(Expr(:macrocall, Expr(:(.), :Threads, QuoteNode(Symbol("@threads"))), __source__, ex))
        else
            $ex
        end
    end |> esc
end