aluo-x

import torch

@torch.jit.script
def distance1(input_pos_1:torch.Tensor, input_pos_2:torch.Tensor) -> torch.Tensor:
    # naive approach
    # n, 3
    # m, 3
    # return n, m
    # n, 1, 3 - 1, m 3
    return torch.sum(torch.square((input_pos_1.unsqueeze(1) - input_pos_2[None])), dim=2)

aluo-x

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP
from fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP
from fairscale.optim.oss import OSS
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
import os

aluo-x

import embree
import numpy as np


def get_centroids(V, F):
    return V[F].mean(axis=1)


def get_cross_products(V, F):
    V0 = V[F][:, 0, :]

aluo-x

elif self.DYN_ARCH in ['single-fc-leaky-wide-discrete']:
    objs_state, agents_state = state
    object_class = extra
    agent_feat = tf.reshape(agents_state, [B, N_AGENT_JOINTS * AGENT_DIM])
    ##### Only allow 1 onj!!!!! #########
    obj_feat = objs_state[:, 0]
    net_input = tf.concat([obj_feat, agent_feat, action], axis=-1)
    backbone_out = self.fcnet(net_input, is_train)

    net_output = self.offsetnet(backbone_out, is_train)