merrymercy

from functools import partial

import numpy as np
import jax
import jax.numpy as jnp

def split(a, axis, factor):
    assert a.shape[axis] % factor == 0
    new_shape = a.shape[:axis] + (factor, a.shape[axis] // factor) + a.shape[axis+1:]
    a = a.reshape(new_shape)

merrymercy

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#

merrymercy

import numpy as np
import tvm
from tvm import te, auto_scheduler, topi

@auto_scheduler.register_workload
def dense_layer(in_dim, out_dim):
    data = te.placeholder((1, in_dim), name="data")
    weight = te.placeholder((out_dim, in_dim), name="weight")
    bias = te.placeholder((out_dim,), name="bias")
    out = topi.nn.dense(data, weight)

merrymercy

struct Edge {
  int src;
  int dst;
  float feature[100];
}

struct Node {
  int node_type;
  int id;

merrymercy

Features

Loop-based Feature

For an IterVar (or an axis), it has three kinds of features

• axis attribute
• arithmetic feature
• touch feature

Axis Attribute