cycyyy

from deepctr_torch.inputs import SparseFeat, DenseFeat
import numpy as np
import torch
from torch import nn
import torch.utils.data as td
import torch.nn.functional as F
from tqdm import tqdm

import sys
MAX = sys.maxsize

cycyyy

from collections import OrderedDict
from tqdm import tqdm
import torch.nn.functional as F
import torch.utils.data as td
from torch import nn
import torch
import numpy as np
from deepctr_torch.inputs import SparseFeat, DenseFeat
import os
import sys

cycyyy

import numpy as np 
from sklearn.neighbors import NearestNeighbors
from scipy import spatial
import math
from multiprocessing.pool import Pool
import multiprocessing

def main(X, Y, k, alpha):
    nbrs = NearestNeighbors(n_neighbors=k, algorithm='kd_tree').fit(X)
    tree = spatial.KDTree(Y)

cycyyy

import numpy as np
from jax import random
from neural_tangents import stax

random_key = random.PRNGKey(42)
SAMPLE_SIZE = 100
BATCH_SIZE = 25


def get_mlp_kernel_fn():