vyeevani

from typing import Sequence
import einops
import jax
import equinox

def make_fourier_features(max_seq_len, embedding_size):
    min_freq = 1.0
    max_resolution = max_seq_len
    num_bands = embedding_size // 2
    freq_bands = jax.numpy.linspace(min_freq, max_resolution / 2, num=num_bands)

vyeevani

use serialport::{self, SerialPort};
use std::time::Duration;
use std::io::{self, Write};

const PING: u8 = 0x01;
const READ: u8 = 0x02;
const WRITE: u8 = 0x03;
const REGWRITE: u8 = 0x04;
const ACTION: u8 = 0x05;
const SYNCWRITE: u8 = 0x83;

vyeevani

import jax
import einops
import equinox

def linear_beta_schedule(num_steps, beta_start, beta_end):
    return jax.numpy.linspace(beta_start, beta_end, num_steps)

def cosine_beta_schedule(num_steps, s=0.008):
    """
    cosine schedule

vyeevani

import jax
import einops
import equinox
from perceiver import Perceiver

class DiffusionPerciever(equinox.Module):
    perceiver: equinox.Module
    input_start: jax.Array
    def __init__(self, context_shapes, input_shape, latent_size, latent_count, key):
        rng = key

vyeevani

import jax
import chex

@chex.dataclass
class OptionalPyTree:
  array: jax.Array
  exists: bool

def make_optional_pytree(array, exists):
    return OptionalPyTree(

vyeevani

import numpy as np
import jax
import equinox
import einops
import typing

def make_attention(query_key_dimension, value_dimension):
    scale = np.sqrt(query_key_dimension)
    def attention(query, key, value, mask):
        """

vyeevani

import os
import glob
import re
import tensorflow as tf
import json
import numpy as np
from PIL import Image

def make_json_parser(json_file, image_dir, image_size, greyscale_images=False):
    def json_parser():

vyeevani

import typing
from operator import mul
from functools import reduce, partial
from collections.abc import Iterable
import jax
import eqxvision as eqv
import equinox as eqx
from eqxvision.utils import CLASSIFICATION_URLS
import equinox
import jax

vyeevani

import typing
from operator import mul
from functools import reduce, partial
from collections.abc import Iterable
import jax
import eqxvision as eqv
import equinox as eqx
from eqxvision.utils import CLASSIFICATION_URLS
import equinox
import jax

vyeevani