text:
import numba
import numpy as np

@numba.jit(nopython=True, cache=True)
def nb_binom(n, k):
    """
    numba version of binomial coefficient function

    Args:
        n (int): how many options
        k (int): how many are chosen
    Returns:
        int: how many ways of choosing
    """
    if k < 0 or k > n:
        return 0
    if k == 0 or k == n:
        return 1
    binom = 1
    for i in range(min(k, n - k)):
        binom *= (n - i)
        binom //= (i + 1)
    return binom

@numba.jit(nopython=True, cache=True)
def precompute_binoms(max_binom):
    """
    precompute binomial coefficients, return as a 2d array

    Args:
        max_binom (int): max value of n in the binomial
    Returns:
        array: max_binom+1 x max_binom+1 array of binomial coefficients
    """
    binoms = np.zeros((max_binom+1, max_binom+1), dtype=type(max_binom))
    for i in range(max_binom+1):
        for j in range(max_binom+1):
            binoms[i,j] = nb_binom(i,j)
    return binoms

@numba.jit(nopython=True, cache=True)
def nb_ix(arr, rows, cols):
    """
    numba implementation of np.ix_
    Args:
        arr (2d array) : matrix to take submatrix of
        rows (array) : rows to be selected in submatrix
        cols (array) : columns to be selected in submatrix
    Return:
        len(rows) * len(cols) array : selected submatrix of arr
    """
    return arr[rows][:, cols]

def matched_reps(reps):
    """
    takes the repeated rows and find a way to pair them up
    to create a perfect matching with many repeated edges

    Args:
        reps (list): list of repeated rows/cols

    Returns:
        x (array): vertex pairs, length 2N for N edges. Index i is matched with i+N
        edgereps (array): length N array for how many times each edge is repeated
        oddmode (int or None): index of odd mode, None if even number of vertices
    """
    n = len(reps)

    if sum(reps) == 0:
        return np.array([], dtype=int), np.array([], dtype=int), None

    # need to pair off the indices with high numbers of repetitions...
    x = range(n) # the starting set of indices
    edgesA = []  # contains part A of each pair
    edgesB = [] # part B of each pair
    edgereps = [] # number of repetitions of a pair
    reps, x = zip(*sorted(zip(reps, x), reverse=True)) # sort according to reps, in descending order
    reps = list(reps)
    x = list(x)

    # remove zeros
    nonzero_reps = []
    nonzero_x = []
    for i, r in zip(x, reps):
        if r > 0:
            nonzero_reps.append(r)
            nonzero_x.append(i)
    reps = nonzero_reps
    x = nonzero_x

    while len(reps) > 1 or (len(reps) == 1 and reps[0] > 1):
        reps, x = zip(*sorted(zip(reps, x), reverse=True)) #sort
        reps = list(reps)
        x = list(x)
        if len(reps) == 1 or reps[0] > reps[1] * 2:
            #if largest number of reps is more than double the 2nd largest, pair it with itself
            edgesA += [x[0]]
            edgesB += [x[0]]
            edgereps += [reps[0]//2]
            if reps[0] % 2 == 0:
                x = x[1:]
                reps = reps[1:]
            else:
                reps[0] = 1
        else:
            #otherwise, form pairs between largest reps and 2nd largest reps
            edgesA += [x[0]]
            edgesB += [x[1]]
            edgereps += [reps[1]]
            if reps[0] > reps[1]:
                if len(x) > 2:
                    x = [x[0]] + x[2:]
                    reps = [reps[0]-reps[1]] + reps[2:]
                else:
                    x = [x[0]]
                    reps = [reps[0]-reps[1]]
            else:
                x = x[2:]
                reps = reps[2:]

    if len(x)==1:
        oddmode=x[0] #if there is an unpaired mode, store it
    else:
        oddmode=None

    # the adjacency matrix of red edges connects 1 to N/2+1, 2 to N/2+2, etc.
    # Reorder the indices (from x2 back to x) so that the paired indices get
    # connected by red edges
    x = np.asarray(edgesA + edgesB, dtype=np.int64) #reordered list of indices
    edgereps = np.asarray(edgereps, dtype=np.int64)

    return x, edgereps, oddmode

@numba.jit(nopython=True, cache=True)
def find_kept_edges(j, reps):
    """
    write j as a string where the ith digit is in base (reps[i]+1)
    decides which edges are included given index of the inclusion/exclusion sum

    Args:
        j (int): index of sum
        reps (list): number of repetitions of each edge
    Returns:
        array : number of repetitions kept for the current inclusion/exclusion step
    """
    num = j
    output = []
    bases = np.asarray(reps) + 1
    for base in bases[::-1]:
        output.append(num % base)
        num //= base
    return np.array(output[::-1], dtype=reps.dtype)

@numba.jit(nopython=True, cache=True)
def f(E, n):
    """
    evaluate the polyonial coefficients of the function in the eigevalue-trace formula

    Args:
        E (array): eigevavlues of AX
        n (int): number of polynomial coefficients to compute
    Returns:
        array: polynomial coefficients
    """
    E_k = E.copy()
    # Compute combinations in O(n^2log n) time
    # code translated from thewalrus matlab script
    count = 0
    comb = np.zeros((2, n//2+1), dtype=np.complex128)
    comb[0,0] = 1
    for i in range(1, n//2+1):
        factor = E_k.sum() / (2 * i)
        E_k *= E
        powfactor = 1
        count = 1 - count
        comb[count, :] = comb[1-count, :]
        for j in range(1, n//(2*i)+1):
            powfactor *= factor / j
            for k in range(i*j+1, n//2+2):
                comb[count, k-1] += comb[1-count, k-i*j-1] * powfactor
    return comb[count, :]

@numba.jit(nopython=True, cache=True)
def f_loop(E, AX_S, XD_S, D_S, n):
    """
    evaluate the polyonial coefficients of the function in the eigevalue-trace formula

    Args:
        E (array): eigevavlues of AX
        AX_S (array): AX_S with weights given by repetitions and exluded rows removed
        XD_S (array): Diagonal multiplied by X
        D_S (array): Diagonal
        n (int): number of polynomial coefficients to compute
    Returns:
        array: polynomial coefficients
    """
    E_k = E.copy()
    # Compute combinations in O(n^2log n) time
    # code translated from thewalrus matlab script
    count = 0
    comb = np.zeros((2, n//2+1), dtype=np.complex128)
    comb[0,0] = 1
    for i in range(1, n//2+1):
        factor = E_k.sum() / (2 * i) + (XD_S @ D_S) / 2
        E_k *= E
        XD_S = XD_S @ AX_S
        powfactor = 1
        count = 1 - count
        comb[count, :] = comb[1-count, :]
        for j in range(1, n//(2*i)+1):
            powfactor *= factor / j
            for k in range(i*j+1, n//2+2):
                comb[count, k-1] += comb[1-count, k-i*j-1] * powfactor
    return comb[count, :]

@numba.jit(nopython=True, cache=True)
def f_loop_odd(E, AX_S, XD_S, D_S, n, oddloop, oddVX_S):
    """
    evaluate the polyonial coefficients of the function in the eigevalue-trace formula
    when there is a self-edge in the fixed perfect matching

    Args:
        E (array): eigevavlues of AX
        AX_S (array): AX_S with weights given by repetitions and exluded rows removed
        XD_S (array): Diagonal multiplied by X
        D_S (array): Diagonal
        n (int): number of polynomial coefficients to compute
        oddloop (float): weight of self-edge
        oddVX_S (array): vector corresponding to matrix at the index of the self-edge
    Returns:
        array: polynomial coefficients
    """
    E_k = E.copy()

    count = 0
    comb = np.zeros((2, n+1), dtype=np.complex128)
    comb[0,0] = 1
    for i in range(1, n+1):
        if i==1:
            factor = oddloop
        elif i%2==0:
            factor = E_k.sum() / i + (XD_S @ D_S) / 2
            E_k *= E
        else:
            factor = oddVX_S @ D_S
            D_S = AX_S @ D_S

        powfactor = 1
        count = 1 - count
        comb[count, :] = comb[1-count, :]
        for j in range(1, n//i+1):
            powfactor *= factor / j
            for k in range(i*j+1, n+2):
                comb[count, k-1] += comb[1-count, k-i*j-1] * powfactor

    return comb[count, :]

@numba.jit(nopython=True, cache=True)
def get_AX_S(kept_edges, A):
    """
    given the kept edges, return the appropriate scaled submatrices to compute f

    Args:
        kept_edges (array): number of repetitions of each edge
        A (array): matrix before repetitions applied
    Returns:
        AX_nonzero (array): scaled A @ X (where X = ((0, I), (I, 0))
    """

    z = np.concatenate((kept_edges, kept_edges))
    nonzero_rows = np.where(z != 0)[0]
    n_nonzero_edges = len(nonzero_rows) // 2

    kept_edges_nonzero = kept_edges[np.where(kept_edges != 0)]

    A_nonzero = nb_ix(A, nonzero_rows, nonzero_rows)

    AX_nonzero = np.empty_like(A_nonzero, dtype=np.complex128)
    AX_nonzero[:,:n_nonzero_edges] = kept_edges_nonzero * A_nonzero[:,n_nonzero_edges:]
    AX_nonzero[:,n_nonzero_edges:] = kept_edges_nonzero * A_nonzero[:,:n_nonzero_edges]

    return AX_nonzero

@numba.jit(nopython=True, cache=True)
def get_submatrices(kept_edges, A, D, oddV):
    """
    given the kept edges, return the appropriate scaled submatrices to compute f

    Args:
        kept_edges (array): number of repetitions of each edge
        A (array): matrix before repetitions applied
        D (array): diagonal before repetitions applied
        oddV (array): row of matrix at index of self-edge, None is no self-edge
    Returns:
        AX_nonzero (array): scaled A @ X (where X = ((0, I), (I, 0))
        XD_nonzero (array): scaled X @ D
        D_nonzero (array): scaled D
        oddVX_nonzero (array): scaled oddV @ X
    """

    z = np.concatenate((kept_edges, kept_edges))
    nonzero_rows = np.where(z != 0)[0]
    n_nonzero_edges = len(nonzero_rows) // 2

    kept_edges_nonzero = kept_edges[np.where(kept_edges != 0)]

    A_nonzero = nb_ix(A, nonzero_rows, nonzero_rows)

    AX_nonzero = np.empty_like(A_nonzero, dtype=np.complex128)
    AX_nonzero[:,:n_nonzero_edges] = kept_edges_nonzero * A_nonzero[:,n_nonzero_edges:]
    AX_nonzero[:,n_nonzero_edges:] = kept_edges_nonzero * A_nonzero[:,:n_nonzero_edges]

    D_nonzero = D[nonzero_rows]

    XD_nonzero = np.empty_like(D_nonzero, dtype=np.complex128)
    XD_nonzero[:n_nonzero_edges] = kept_edges_nonzero * D_nonzero[n_nonzero_edges:]
    XD_nonzero[n_nonzero_edges:] = kept_edges_nonzero * D_nonzero[:n_nonzero_edges]

    if oddV is not None:
        oddV_nonzero = oddV[nonzero_rows]
        oddVX_nonzero = np.empty_like(oddV_nonzero, dtype=np.complex128)
        oddVX_nonzero[:n_nonzero_edges] = kept_edges_nonzero * oddV_nonzero[n_nonzero_edges:]
        oddVX_nonzero[n_nonzero_edges:] = kept_edges_nonzero * oddV_nonzero[:n_nonzero_edges]
    else:
        oddVX_nonzero = None

    return AX_nonzero, XD_nonzero, D_nonzero, oddVX_nonzero

@numba.jit(nopython=True, cache=True)
def get_submatrix_batch_odd0(kept_edges, oddV0):
    """
    find oddVX_nonzero0 for batching (sometimes different vertices are identified as self edges)

    Args:
        kept_edges (array): number of repetitions of each edge
        oddV0 (array): row of matrix at index of self-edge, None is no self-edge

    Returns:
        array: scaled oddV0 @ X for

    """
    z = np.concatenate((kept_edges, kept_edges))
    nonzero_rows = np.where(z != 0)[0]
    n_nonzero_edges = len(nonzero_rows) // 2

    kept_edges_nonzero = kept_edges[np.where(kept_edges != 0)]
    oddV_nonzero0 = oddV0[nonzero_rows]
    oddVX_nonzero0 = np.empty_like(oddV_nonzero0, dtype=np.complex128)
    oddVX_nonzero0[:n_nonzero_edges] = kept_edges_nonzero * oddV_nonzero0[n_nonzero_edges:]
    oddVX_nonzero0[n_nonzero_edges:] = kept_edges_nonzero * oddV_nonzero0[:n_nonzero_edges]

    return oddVX_nonzero0

@numba.jit(nopython=True, cache=True)
def get_Dsubmatrices(kept_edges, D):
    """
    find submatrices for batch gamma functions
    """

    z = np.concatenate((kept_edges, kept_edges))
    nonzero_rows = np.where(z != 0)[0]
    n_nonzero_edges = len(nonzero_rows) // 2

    kept_edges_nonzero = kept_edges[np.where(kept_edges != 0)]

    D_nonzero = D[nonzero_rows]

    XD_nonzero = np.empty_like(D_nonzero, dtype=np.complex128)
    XD_nonzero[:n_nonzero_edges] = kept_edges_nonzero * D_nonzero[n_nonzero_edges:]
    XD_nonzero[n_nonzero_edges:] = kept_edges_nonzero * D_nonzero[:n_nonzero_edges]

    return XD_nonzero, D_nonzero

@numba.jit(nopython=True, cache=True)
def eigvals(M):
    return np.linalg.eigvals(M)

# @numba.jit(nopython=True, cache=True)
def calc_approx_steps(fixed_reps, N_cutoff):
    steps = int(np.prod(np.sqrt(fixed_reps))+1) * N_cutoff // 2
    return steps

max_stars_repo_name: apprenticeadi/gbs

max_issues_repo_name: apprenticeadi/gbs

max_forks_repo_name: apprenticeadi/gbs

reason: import numpy,import numba


text:
"""Testing for discretizers."""

import numpy as np
import pytest
import re
from pyts.preprocessing.discretizer import _uniform_bins, _digitize
from pyts.preprocessing import KBinsDiscretizer

n_samples, n_timestamps = 3, 10
X = np.arange(n_samples * n_timestamps).reshape(n_samples, n_timestamps)

@pytest.mark.parametrize(
    'sample_min, sample_max, n_bins, arr_desired',
    [(np.arange(50), np.arange(20, 70), 2, [np.arange(10, 60)]),
     (np.arange(30, 80), np.arange(-10, 40), 2, [np.arange(10, 60)]),
     (np.arange(30), np.arange(30, 60), 3,
      [np.arange(10, 40), np.arange(20, 50)])]
)
def test_uniform_bins(sample_min, sample_max, n_bins, arr_desired):
    """Test '_uniform_bins' function."""
    arr_actual = _uniform_bins(sample_min, sample_max,
                               n_samples=sample_min.size, n_bins=n_bins)
    np.testing.assert_array_equal(arr_actual, arr_desired)

@pytest.mark.parametrize(
    'X, bins, arr_desired',
    [(X, np.array([4.5, 14.5, 24.5]),
      [[i] * 5 + [i + 1] * 5 for i in range(3)]),
     (X, np.array([[-0.5, 1.5, 8.5, 11.5],
                   [9.5, 11.5, 18.5, 21.5],
                   [19.5, 21.5, 28.5, 31.5]]),
      [[1] * 2 + [2] * 7 + [3] for _ in range(3)])]
)
def test_digitize(X, bins, arr_desired):
    """Test '_digitize' function."""
    arr_actual = _digitize(X, bins)
    np.testing.assert_array_equal(arr_actual, arr_desired)

@pytest.mark.parametrize(
    'params, error, err_msg',
    [({'n_bins': None}, TypeError, "'n_bins' must be an integer."),

     ({'n_bins': 1}, ValueError,
      "'n_bins' must be greater than or equal to 2 and lower than "
      "or equal to n_timestamps (got 1)."),

     ({'n_bins': n_timestamps + 1}, ValueError,
      "'n_bins' must be greater than or equal to 2 and lower than "
      "or equal to n_timestamps (got {0}).".format(n_timestamps + 1)),

     ({'strategy': 'whoops'}, ValueError,
      "'strategy' must be either 'uniform', 'quantile' or 'normal' (got {0})."
      .format('whoops'))]
)
def test_parameter_check(params, error, err_msg):
    """Test parameter validation."""
    discretizer = KBinsDiscretizer(**params)
    with pytest.raises(error, match=re.escape(err_msg)):
        discretizer.transform(X)

def test_constant_sample():
    """Test that a ValueError is raised with a constant sample."""
    discretizer = KBinsDiscretizer()
    with pytest.raises(ValueError, match="At least one sample is constant."):
        discretizer.fit_transform(np.ones((10, 15)))

def test_warning_smaller_n_bins():
    """Test that a warning is raised when n_bins will be smaller."""
    X_new = np.r_[np.zeros((1, n_timestamps)), X]
    discretizer = KBinsDiscretizer()
    warning_msg = (
        "Some quantiles are equal. The number of bins will be "
        "smaller for sample {0}. Consider decreasing the number "
        "of bins or removing these samples.".format([0])
    )
    with pytest.warns(None, match=re.escape(warning_msg)):
        discretizer._compute_bins(
            X_new, n_samples + 1, n_bins=5, strategy='quantile'
        )

@pytest.mark.parametrize(
    'params, X, arr_desired',
    [({'n_bins': 10, 'strategy': 'quantile'}, X,
      [list(range(10))] * n_samples),

     ({'n_bins': 10, 'strategy': 'uniform'}, X,
      [list(range(10))] * n_samples),

     ({'n_bins': 5, 'strategy': 'quantile'}, X,
      [np.repeat(np.arange(5), 2)] * n_samples),

     ({'n_bins': 5, 'strategy': 'uniform'}, X,
      [np.repeat(np.arange(5), 2)] * n_samples),

     ({'n_bins': 3, 'strategy': 'normal'}, [[-0.5, 0, 0.5]], [[0, 1, 2]])]
)
def test_actual_results(params, X, arr_desired):
    """Test that the actual results are the expected ones."""
    arr_actual = KBinsDiscretizer(**params).fit_transform(X)
    np.testing.assert_allclose(arr_actual, arr_desired, atol=1e-5, rtol=0.)

ext: py

max_stars_repo_name: martanto/pyts

max_issues_repo_name: martanto/pyts

max_forks_repo_name: martanto/pyts

reason: import numpy


text:
#!/usr/bin/env python

import rospy
import actionlib
from move_base_msgs.msg import MoveBaseAction, MoveBaseGoal
import numpy as np

class SimpleNavigationGoals:

    def __init__(self):
        pass

    def go_to(self, x, y, theta):
        client = None
        client = actionlib.SimpleActionClient('move_base',MoveBaseAction)
        client.wait_for_server()

        goal = MoveBaseGoal()
        goal.target_pose.header.frame_id = "map"
        goal.target_pose.header.stamp = rospy.Time.now()

        theta = self.euler_to_quaternion(theta)

        goal.target_pose.pose.position.x = x
        goal.target_pose.pose.position.y = y
        goal.target_pose.pose.position.z = 0
        goal.target_pose.pose.orientation.w = theta

        client.send_goal(goal)
        wait = client.wait_for_result()

        if not wait:
            rospy.logerr("Action server not available!")
            rospy.signal_shutdown("Action server not available!")
        else:
            rospy.loginfo("Done")
            return client.get_result()

    def euler_to_quaternion(self, yaw):
        pitch = 0
        roll = 0
        qw = np.cos(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)
        return qw

    def _shutdown(self):
        rospy.signal_shutdown("Shutdown!")


text:
import numpy as np
import trimesh

class ErrorComputingRayTriangleIntersection(Exception):
    """Error when computing the intersection between a ray and a mesh triangle"""
    pass

def mesh_to_grasps(mesh, gripper, MAX_MESH_SIDE_LENGTH=0.2, TOTAL_RAYS_TO_USE=512, use_only_external_points=True):
    # Version 1: Use the face normals for grasp axes
    #            Face normals start at the centroid of their face triangles.
    # Formula: https://en.wikipedia.org/wiki/Centroid#Of_a_triangle
    # mesh_triangles = [[mesh.vertices[triangle[0]], mesh.vertices[triangle[1]], mesh.vertices[triangle[2]]] for triangle in mesh.faces]
    mesh_triangles = np.array(mesh.vertices[mesh.faces]) # Create triplets of triangle points
    face_centroids = np.array([np.mean(triangle, axis=0) for triangle in mesh_triangles])
    # http://jwilson.coe.uga.edu/EMAT6680/Dunbar/Assignment4/Assignment4_KD.htm

    ### Compute the rays: origins and directions (using face normals)
    ray_origins = []
    ray_directions = []

    # Limit to a random TOTAL_RAYS_TO_USE rays
    randomSelectionIndexes = np.arange(len(mesh.face_normals)) # ray_directions = normals
    np.random.shuffle(randomSelectionIndexes)
    # randomSelectionIndexes = randomSelectionIndexes[:TOTAL_RAYS_TO_USE] #
    randomSelectionIndexes = randomSelectionIndexes[:min(TOTAL_RAYS_TO_USE, len(mesh.face_normals))]

    # print(randomSelectionIndexes)

    # print("Ray origins #: {0}".format(ray_origins))
    # print("Ray directions #: {0}".format(ray_directions))
    ray_origins = np.array(face_centroids[randomSelectionIndexes])
    ray_directions = np.array(mesh.face_normals)[randomSelectionIndexes] # Use face normals

    # Face normals have norm 1.0
    # The bounding box of the mesh has side of 0.2 (as defined in the mesh generation file)
    ray_origins = ray_origins - (MAX_MESH_SIDE_LENGTH*ray_directions)
    ray_directions *= 2 * (MAX_MESH_SIDE_LENGTH)

    # Compute intersection points for all these rays
    successfully_computed_intersection_points = False
    while (not successfully_computed_intersection_points):
        try:
            # run the mesh-ray query
            # Returns:
            #     locations ((n) sequence of (m,3) float) - Intersection points
            #     index_ray ((n,) int) - Array of ray indexes
            #     index_tri ((n,) int) - Array of triangle (face) indexes
            locations, index_ray, index_tri = mesh.ray.intersects_location(
                                                    ray_origins=ray_origins,
                                                    ray_directions=ray_directions) # multiple_hits=False)
            successfully_computed_intersection_points = True
        except:
            print("Mesh faces: {}".format(len(mesh.face_normals)))
            # import pdb; pdb.set_trace()
            raise ErrorComputingRayTriangleIntersection

            # print("Retrying computing grasps with fewer rays")
            # successfully_computed_grasps = False
            #     # This solution may fail (if the bad pair of points is at the beginning of the array)
            #     if (len(randomSelectionIndexes) > 32):
            #         randomSelectionIndexes = randomSelectionIndexes[0:(len(randomSelectionIndexes)//2)]
            #         #
            #         ray_origins = np.array(face_centroids[randomSelectionIndexes])
            #         ray_directions = np.array(mesh.face_normals)[randomSelectionIndexes] # Use face normals
            #         #
            #         # Face normals have norm 1.0
            #         ray_origins = ray_origins - (MAX_MESH_SIDE_LENGTH*ray_directions)
            #         ray_directions *= 2 * (MAX_MESH_SIDE_LENGTH)
            #     # else: any other way to solve it?

    # Onwards, compute the Grasp structures:
    #     - group the contact points by rays
    #     - contact points
    #     - faces (IDs) on which these points are located

    # print("Total intersections: {}".format(len(locations)))
    # print("Intersection locations:\n{}".format(locations))
    # print("Ray IDs:\n{}".format(index_ray))
    # print("Triangle IDs:\n{}".format(index_tri))

    # Group the intersection points by the rays that generated them
    # We want to create a structure here of the type:
    # Ray ID: array [intersection points]
    rays_intersection_points = {}
    rays_intersection_faces = {}
    for i in range(len(index_ray)):
        if index_ray[i] not in rays_intersection_points:
            # Create a new entry for this ray
            rays_intersection_points[index_ray[i]] = []
            rays_intersection_faces[index_ray[i]] = []
            # Add the corresponding intersection point
        rays_intersection_points[index_ray[i]].append(locations[i])
        rays_intersection_faces[index_ray[i]].append(index_tri[i])

    #     ### Visualise the rays
    #     # stack rays into line segments for visualization as Path3D
    #     ray_visualize = trimesh.load_path(np.hstack((ray_origins,
    #                                                  ray_origins+ray_directions)).reshape(-1, 2, 3))

    #     # Visualise only the mesh:
    #     # mesh.show()

    #     # unmerge so viewer doesn't smooth
    #     mesh.unmerge_vertices()
    #     # make mesh white-ish
    #     mesh.visual.face_colors = [255,255,255,255]
    #     mesh.visual.face_colors[index_tri] = [255, 0, 0, 255]

    #     # Create a visualization scene with rays, hits, and mesh
    #     #scene = trimesh.Scene([mesh, ray_visualize, mesh.bounding_sphere])
    #     scene = trimesh.Scene([mesh, ray_visualize])

    #     # show the visualization
    #     scene.show()

    # CHECK WHICH INTERSECTION POINTS ARE CLOSER THAN THE GRIPPER WIDTH
    # The container that will hold grasps in format:
    # rayID: (point1, point2)
    grasps = []
    total_grasps=0

    for rayId in rays_intersection_points: # loop through the keys of the dictionary
        # print("\nRay {}:".format(i))
        # for j in range(len(rays_intersections[i])):
        #     print(rays_intersections[i][j])

        # Assumption: each ray intersects the mesh in at least 2 points (false when it touches a face)
        # Try: (max width)
        # If that is wider than maxGripperWidth,
        # then try smaller and smaller until first success
        # or until we run out of points

        # Constraint:
        # Distance between two successive points in empty space
        # should be bigger than finger thickness (X&Y)

        totalRayIntersectionPoints = len(rays_intersection_points[rayId])
        # print("Total intersection points: {}".format(totalRayIntersectionPoints))

        # Ensure you have an even number of intersection points (not true for objects with 0-volume surfaces)
        # assert (totalRayIntersectionPoints%2==0),"Ray #{} has an odd number of intersections: {}".format(rayId,totalRayIntersectionPoints)

        # If you want to use only the external points (to avoid problems with grasp reachability)
        if (use_only_external_points):
            startID = 0
            endID = totalRayIntersectionPoints - 1
            distance = np.linalg.norm(rays_intersection_points[rayId][startID] - rays_intersection_points[rayId][endID])
            # print("Distance: {}".format(distance))
            # Reject pair of points if they are too far apart
            # and if there is only 1 intersection point (i.e. P1 = P2)
            if ((distance < gripper['width']) and (totalRayIntersectionPoints > 1)):
                grasp = {}
                grasp["point_locations"] = (rays_intersection_points[rayId][startID], rays_intersection_points[rayId][endID])
                grasp["point_faces"] = (rays_intersection_faces[rayId][startID], rays_intersection_faces[rayId][endID])
                #
                # Additional info (just in case)
                grasp["ray_id"] = rayId
                grasp["grasp_id"] = total_grasps
                # grasps.append((rays_intersections[rayId][startID],rays_intersections[rayId][endID]))
                grasps.append(grasp)
                total_grasps += 1

        # If you want to
        # Generate grasps at all point pairs on a ray intersecting a mesh (event interior ones)
        else: #if (not use_only_external_points):
            # WARNING: If you do not want to work with pairs of contact points,
            # then you have to check the presence of intersections between the two contact points and the mesh.
            for startID in range(0,totalRayIntersectionPoints,2):
                for endID in range(startID+1,totalRayIntersectionPoints,2):
                    # Test if points are closer than max gripper width
                    # print(rays_intersections[i][startID])
                    # print(rays_intersections[i][endID])
                    distance = np.linalg.norm(rays_intersection_points[rayId][startID] - rays_intersection_points[rayId][endID])
                    # print("Distance: {}".format(distance))
                    if (distance < gripper['width']):
                        grasp = {}
                        grasp["point_locations"] = (rays_intersection_points[rayId][startID], rays_intersection_points[rayId][endID])
                        grasp["point_faces"] = (rays_intersection_faces[rayId][startID], rays_intersection_faces[rayId][endID])
                        #
                        # Additional info (just in case)
                        grasp["ray_id"] = rayId
                        grasp["grasp_id"] = total_grasps
                        # grasps.append((rays_intersections[rayId][startID],rays_intersections[rayId][endID]))
                        grasps.append(grasp)
                        total_grasps += 1

    # Check
    print("Total rays: {}".format(len(rays_intersection_points)))
    print("Total grasps: {} ({})".format(len(grasps), total_grasps))
    print("Mesh volume: {}".format(mesh.volume))
    # print(grasps)
    if (total_grasps == 0):
        # Debugging
        mesh.show()

    return grasps


text:
# install rdkit library on colab environment
!wget -c https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh
!chmod +x Miniconda3-latest-Linux-x86_64.sh
!time bash ./Miniconda3-latest-Linux-x86_64.sh -b -f -p /usr/local
!time conda install -q -y -c conda-forge rdkit

# %matplotlib inline
import matplotlib.pyplot as plt
import sys
import os
sys.path.append('/usr/local/lib/python3.7/site-packages/')

#import basic python packages
import numpy as np
import pandas as pd

from rdkit import Chem
from rdkit.Chem import DataStructs
from rdkit.Chem import AllChem
from rdkit.Chem import RDConfig
from rdkit import rdBase

import tensorflow as tf
import tensorflow.keras as keras
from keras import Model
from keras.backend import one_hot
from keras.layers import Activation, Dense, Dropout, Input
from keras.utils import np_utils

# load data from files
from google.colab import files
data_AR = files.upload() # select 'nr-ar.smiles' file and upload it
data_ER = files.upload() # select 'nr-er.smiels' file and upload it

# save the uploaded data as Pandas dataframe
import io
df_AR = pd.read_csv(io.BytesIO(data_AR['nr-ar.smiles']), sep = '\t', names = ['smiles', 'array', 'activity'])
df_ER = pd.read_csv(io.BytesIO(data_ER['nr-er.smiles']), sep = '\t', names = ['smiles', 'array', 'activity'])

smi_AR = df_AR['smiles']
smi_ER = df_ER['smiles']
activity_AR = df_AR['activity']
activity_ER = df_ER['activity']

# data preprocessing; convert SMILES string to 2048 bit Fingerprint
def preprocessing(data_X, data_Y):
###################
# Your code here
###################
  return np.array(data_x), np.array(data_y)

data_X_AR, data_Y_AR = preprocessing(smi_AR, activity_AR)
data_X_ER, data_Y_ER = preprocessing(smi_ER, activity_ER)

# check the shape of the raw data
print(data_X_AR.shape)
print(data_Y_AR.shape)
print(data_X_ER.shape)
print(data_Y_ER.shape)

# split the training, validation and test dataset. train:validation:test=8:1:1
from sklearn.model_selection import train_test_split
#split training and test dataset from raw dataset
###################
# Your code here
###################

#split training and validation dataset from training dataset
###################
# Your code here
###################

# check the shape of the splited data
print(X_train_AR.shape, Y_train_AR.shape)
print(X_valid_AR.shape, Y_valid_AR.shape)
print(X_test_AR.shape, Y_test_AR.shape)
print(X_train_ER.shape, Y_train_ER.shape)
print(X_valid_ER.shape, Y_valid_ER.shape)
print(X_test_ER.shape, Y_test_ER.shape)

# build a neural network model
# the activation function of hidden layer is relu, and output is softmax

###################
# Your code here
###################

# make a prediction model for AR
###################
# Your code here
###################
model_AR.summary() # print summary of AR model

# make a predictin model for ER
###################
# Your code here
###################
model_ER.summary()# print summary of ER model

# early stopping with validation dataset
from keras.callbacks import EarlyStopping
from keras.callbacks import ModelCheckpoint

es_AR = EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=5)
mc_AR = ModelCheckpoint('best_model_AR.h5', monitor='val_acc', mode='max', verbose=1, save_best_only=True)
es_ER = EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=5)
mc_ER = ModelCheckpoint('best_model_ER.h5', monitor='val_acc', mode='max', verbose=1, save_best_only=True)

# training the model using training and validation set
epochs = 50
hist_AR = model_AR.fit(X_train_AR, Y_train_AR, validation_data=(X_valid_AR, Y_valid_AR), batch_size=32, epochs=epochs, verbose=0, callbacks=[es_AR, mc_AR])
hist_ER = model_ER.fit(X_train_ER, Y_train_ER, validation_data=(X_valid_ER, Y_valid_ER),batch_size=32, epochs=epochs, verbose=0, callbacks=[es_ER, mc_ER])

# plot the accuracy and loss value of each model.
###################
# Your code here
###################

# evaluate the prediction model using test dataset.
from sklearn.metrics import classification_report, roc_curve, auc
# AUC of AR dataset
###################
# Your code here
###################

# AUC of ER dataset
###################
# Your code here
###################


text:
#%%
import os
import glob
import numpy as np
import scipy as sp
import pandas as pd
import re
import git

# Import matplotlib stuff for plotting
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib as mpl

# Seaborn, useful for graphics
import seaborn as sns

# Import the project utils
import ccutils

# Set PBoC plotting format
ccutils.viz.set_plotting_style()

#%%

# Find home directory for repo
repo = git.Repo("./", search_parent_directories=True)
homedir = repo.working_dir

# Define directories for data and figure
figdir = f'{homedir}/fig/si/'
datadir = f'{homedir}/data/csv_gillespie/'

# %%
df_sim_prot = pd.read_csv(datadir + "two_state_protein_gillespie.csv")

# Extract mRNA data
mRNA_names = [x for x in df_sim_prot.columns if re.match(r"[m]\d", x)]
mRNA_data = df_sim_prot.loc[:, mRNA_names].values
# Compute mean mRNA
mRNA_mean = mRNA_data.mean(axis=1)

# Extract protein data
protein_names = [x for x in df_sim_prot.columns if re.match(r"[p]\d", x)]
protein_data = df_sim_prot.loc[:, protein_names].values
# Compute mean protein
protein_mean = protein_data.mean(axis=1)

# Initialize plot
fig, ax = plt.subplots(2, 1, figsize=(2.5, 2), sharex=True)

# Define colors
colors = sns.color_palette("Paired", n_colors=2)

# Define time stepsize for plot
binstep = 10
# Define every how many trajectories to plot
simnum = 10

# Plot mRNA trajectories
ax[0].plot(
    df_sim_prot["time"][0::binstep] / 60,
    mRNA_data[0::binstep, 0::simnum],
    color=colors[0],
)
# Plot mean mRNA
ax[0].plot(
    df_sim_prot["time"][0::binstep] / 60,
    mRNA_mean[0::binstep],
    color=colors[1],
)

# Plot protein trajectories
ax[1].plot(
    df_sim_prot["time"][0::binstep] / 60,
    protein_data[0::binstep, 0::simnum],
    color=colors[0],
)
# Plot mean protein
ax[1].plot(
    df_sim_prot["time"][0::binstep] / 60,
    protein_mean[0::binstep],
    color=colors[1],
)

# Group data frame by cell cycle
df_group = df_sim_prot.groupby("cycle")
# Loop through cycles
for i, (group, data) in enumerate(df_group):
    # Define the label only for the last cell cycle not to repeat in legend
    if group == df_sim_prot["cycle"].max():
        label_s = "single promoter"
        label_d = "two promoters"
    else:
        label_s = ""
        label_d = ""
    # Find index for one-promoter state
    idx = np.where(data.state == "single")[0]
    # Indicate states with two promoters
    ax[0].axvspan(
        data.iloc[idx.min()]["time"] / 60,
        data.iloc[idx.max()]["time"] / 60,
        facecolor="#e3dcd1",
        label=label_s,
    )
    ax[1].axvspan(
        data.iloc[idx.min()]["time"] / 60,
        data.iloc[idx.max()]["time"] / 60,
        facecolor="#e3dcd1",
        label=label_s,
    )

    # Find index for two-promoter state
    idx = np.where(data.state == "double")[0]
    # Indicate states with two promoters
    ax[0].axvspan(
        data.iloc[idx.min()]["time"] / 60,
        data.iloc[idx.max()]["time"] / 60,
        facecolor="#ffedce",
        label=label_d,
    )
    ax[1].axvspan(
        data.iloc[idx.min()]["time"] / 60,
        data.iloc[idx.max()]["time"] / 60,
        facecolor="#ffedce",
        label=label_d,
    )

# Set limits
ax[0].set_xlim(df_sim_prot["time"].min() / 60, df_sim_prot["time"].max() / 60)
# Label plot
ax[1].set_xlabel("time (min)")
ax[0].set_ylabel("mRNA/cell")
ax[1].set_ylabel("protein/cell")

# Set legend for both plots
ax[0].legend(
    loc="upper left",
    ncol=2,
    frameon=False,
    bbox_to_anchor=(-0.12, 0, 0, 1.3),
    fontsize=6.5,
)

# Align y axis labels
fig.align_ylabels()

plt.subplots_adjust(hspace=0.05)

plt.savefig(figdir + "figS22.pdf", bbox_inches="tight")

reason: import numpy,import scipy


text:
# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import gc

import numpy as np
import pytest

import megengine as mge
import megengine.distributed as dist
import megengine.functional as F
from megengine.core._imperative_rt import CompNode, TensorAttr, imperative
from megengine.core._imperative_rt.core2 import TensorWeakRef, apply, sync
from megengine.core.autodiff.grad import Grad
from megengine.core.ops.builtin import Elemwise, Identity
from megengine.functional.distributed import remote_recv, remote_send

def _elwise(mode):
    op = Elemwise(mode)

    def f(*args):
        (result,) = apply(op, *args)
        return result

    return f

add = _elwise(Elemwise.Mode.ADD)
mul = _elwise(Elemwise.Mode.MUL)
cos = _elwise(Elemwise.Mode.COS)
relu = _elwise(Elemwise.Mode.RELU)

def as_tensor(x):
    return mge.Tensor(x)

def save_to(self, name="grad"):
    def callback(grad):
        setattr(self, name, grad)

    return callback

@pytest.mark.require_ngpu(2)
@pytest.mark.isolated_distributed
def test_dist_grad():
    world_size = 2
    x_np = np.random.rand(10).astype("float32")

    @dist.launcher
    def worker():
        rank = dist.get_rank()
        if rank == 0:
            with Grad() as grad:
                x = as_tensor(x_np)
                grad.wrt(x, callback=save_to(x))
                # need a placeholder to trace operator
                remote_send(x, 1)
                recv_x = remote_recv(1)
                y = recv_x * recv_x
                grad([y], [as_tensor(np.ones_like(x_np))])
            np.testing.assert_almost_equal(x.grad.numpy(), x.numpy() * 2)
        elif rank == 1:
            with Grad() as grad:
                recv_x = remote_recv(0)
                remote_send(recv_x, 0)
                grad([], [])

    worker()

def test_grad():
    x_np = np.random.rand(10).astype("float32")
    x = as_tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        y = cos(x)
        grad(y, as_tensor(np.ones_like(x_np)))

    np.testing.assert_almost_equal(x.grad.numpy(), -np.sin(x_np))

def test_grad_2():
    x_np = np.random.rand(10).astype("float32")
    x = as_tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        y = mul(x, x)
        y = mul(y, y)
        grad(y, as_tensor(np.ones_like(x_np)))

    np.testing.assert_almost_equal(x.grad.numpy(), 4 * x_np ** 3, decimal=6)

@pytest.mark.require_higher_order_directive()
def test_2nd_grad():
    x_np = np.random.rand(10).astype("float32")
    x = as_tensor(x_np)
    ones = as_tensor(np.ones_like(x_np))

    with Grad("grad2") as grad2:
        with Grad("grad") as grad:
            grad2.wrt(x, callback=save_to(x))
            grad.wrt(x, callback=save_to(x))
            y = cos(x)
            grad(y, ones)
            z = x.grad
            np.testing.assert_almost_equal(x.grad.numpy(), -np.sin(x_np), decimal=5)

        x.grad = None
        grad2(z, ones)

        np.testing.assert_almost_equal(x.grad.numpy(), -np.cos(x_np), decimal=5)

def test_grad_with_tensor_wrapper():
    x_np = np.random.rand(10).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        y = mul(x, x)
        y = mul(y, y)
        grad(y, mge.Tensor(np.ones_like(x_np)))

    np.testing.assert_almost_equal(x.grad.numpy(), 4 * x_np ** 3, decimal=6)

def test_release():
    def check(f):
        n = 0
        d = None
        gc.disable()
        try:
            for i in range(3):
                f()
                m = len(gc.get_objects())
                d = m - n
                n = m
            assert d == 0
        finally:
            gc.enable()

    x = mge.Tensor([0.0])
    dy = mge.Tensor(np.ones_like(x.numpy()))

    @check
    def _():
        with Grad() as g:
            g.wrt(x)
            y = x * x
            g(y, dy)

    @check
    def _():
        with Grad() as g:
            g.wrt(x)
            pass

    @check
    def _():
        with Grad() as g:
            g.wrt(x)
            y = x * x

def test_grad_inplace():
    x_np = np.random.rand(10).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        y = mul(x, x)
        y *= y
        grad(y, mge.Tensor(np.ones_like(x_np)))

    np.testing.assert_almost_equal(x.grad.numpy(), 4 * x_np ** 3, decimal=6)

def test_identity():
    x_np = np.random.rand(10).astype("float32")
    x = mge.Tensor(x_np)
    dy_np = np.random.rand(*x.shape).astype("float32")
    dy = mge.Tensor(dy_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        (y,) = apply(Identity(), x)
        grad(y, dy)

    np.testing.assert_array_equal(x.grad.numpy(), dy_np)

def test_elemwise_add():
    x_np = np.random.rand(10).astype("float32")
    y_np = np.random.rand(10, 10).astype("float32")
    dz_np = np.random.rand(10, 10).astype("float32")
    x = mge.Tensor(x_np)
    y = mge.Tensor(y_np)
    dz = mge.Tensor(dz_np)

    refs = {}

    def f(x, y):
        x = x * 2
        refs["x"] = TensorWeakRef(x)
        refs["y"] = TensorWeakRef(y)
        return x + y

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        z = f(x, y)
        del y
        for k, r in refs.items():
            assert r() is None
        grad(z, dz)

    np.testing.assert_almost_equal(x.grad.numpy(), dz_np.sum(0) * 2, decimal=5)

def test_elemwise_relu():
    x_np = [1.0, -1.0]
    dz_np = [1.0]
    x = mge.Tensor(x_np)
    dz = mge.Tensor(dz_np)

    refs = {}

    def f(x):
        x = x * 2
        refs["x"] = TensorWeakRef(x)
        return relu(x)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        z = f(x)
        assert refs["x"]() is None
        grad(z, dz)

    np.testing.assert_almost_equal(x.grad.numpy(), [2.0, 0])

def test_elemwise_relu_backward_fn():
    op = Elemwise(Elemwise.Mode.RELU)
    attr = TensorAttr()
    attr.dtype = "float32"
    attr.comp_node = "xpux"
    result = imperative.make_backward_graph(op, [attr], [True], [True])
    backward_graph, save_for_backward_mask, input_has_grad = result
    assert save_for_backward_mask == [False, True, True], save_for_backward_mask

def test_reshape():
    x_np = np.random.rand(2, 5).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        refs = {}

        def f(x):
            x = x * 1
            y = x.reshape(5, 2)
            refs["x"] = TensorWeakRef(x)
            return y

        y = f(x)
        for _, r in refs.items():
            assert r() is None
        grad(y, F.ones_like(y))

    np.testing.assert_equal(np.ones((2, 5), dtype=np.float32), x.grad.numpy())

def test_subtensor():
    x_np = np.random.rand(3, 3).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        refs = {}

        def f(x):
            x = x * 1
            y = x[1:-1, :2]
            refs["x"] = TensorWeakRef(x)
            return y

        y = f(x)
        for _, r in refs.items():
            assert r() is None
        grad(y, F.ones_like(y))

    np.testing.assert_equal(
        np.array([[0, 0, 0], [1, 1, 0], [0, 0, 0]], dtype=np.float32), x.grad.numpy()
    )

def test_IndexingMultiAxisVec():
    x_np = np.random.rand(3, 3).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        refs = {}

        def f(x):
            x = x * 1
            y = x[[0, 2], [0, 2]]
            refs["x"] = TensorWeakRef(x)
            return y

        y = f(x)
        for _, r in refs.items():
            assert r() is None
        grad(y, F.ones_like(y))

    np.testing.assert_equal(
        np.array([[1, 0, 0], [0, 0, 0], [0, 0, 1]], dtype=np.float32), x.grad.numpy()
    )

def test_AxisAddRemove():
    x_np = np.random.rand(1, 5).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        refs = {}

        def f(x):
            x = x * 1
            y = F.squeeze(F.expand_dims(x, 2), 0)
            refs["x"] = TensorWeakRef(x)
            return y

        y = f(x)
        for _, r in refs.items():
            assert r() is None
        grad(y, F.ones_like(y))

    np.testing.assert_equal(
        np.array([[1, 1, 1, 1, 1]], dtype=np.float32), x.grad.numpy()
    )

def test_Broadcast():
    x_np = np.random.rand(3, 3, 1).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        y = F.broadcast_to(x, (3, 3, 10))
        grad(y, F.ones_like(y))

    np.testing.assert_equal(np.ones((3, 3, 1), dtype=np.float32) * 10, x.grad.numpy())

def test_interpolate_fastpath():
    x_np = np.random.rand(3, 3, 32, 32).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        y = F.vision.interpolate(x, size=(16, 16), mode="bilinear")
        grad(y, F.ones_like(y))

    np.testing.assert_equal(np.ones(x_np.shape, dtype=np.float32) / 4, x.grad.numpy())

def test_Reduce_sum():
    x_np = np.random.rand(3, 3).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        y = x.sum(axis=0)
        grad(y, F.ones_like(y))

    np.testing.assert_equal(np.ones((3, 3), dtype=np.float32), x.grad.numpy())

def test_Reduce_mean():
    x_np = np.random.rand(3, 3).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        y = x.mean(axis=0)
        grad(y, F.ones_like(y))

    np.testing.assert_equal(np.ones((3, 3), dtype=np.float32) / 3, x.grad.numpy())

def test_addAxis():
    x_np = np.random.rand(3, 3).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        refs = {}

        def f(x):
            x = x * 1
            y = F.expand_dims(x, [2, 3])
            refs["x"] = TensorWeakRef(x)
            return y

        y = f(x)
        for _, r in refs.items():
            assert r() is None
        grad(y, F.ones_like(y))

    np.testing.assert_equal(np.ones((3, 3), dtype=np.float32), x.grad.numpy())

def test_removeAxis():
    x_np = np.random.rand(3, 3, 1, 1).astype("float32")
    x = mge.Tensor(x_np)

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))
        refs = {}

        def f(x):
            x = x * 1
            y = F.squeeze(x, [2, 3])
            refs["x"] = TensorWeakRef(x)
            return y

        y = f(x)
        for _, r in refs.items():
            assert r() is None
        grad(y, F.ones_like(y))

    np.testing.assert_equal(np.ones((3, 3, 1, 1), dtype=np.float32), x.grad.numpy())

def test_dot():
    x = np.random.rand(2, 2).astype("float32")
    x = mge.Tensor(x)
    u = F.ones((2,))
    v = F.ones((2,))

    with Grad() as grad:
        grad.wrt(x, callback=save_to(x))

        def f(x):
            return F.dot(u, F.matmul(x, v))

        y = f(x)
        grad(y, F.ones_like(y))

    np.testing.assert_equal(np.ones((2, 2), dtype=np.float32), x.grad.numpy())

max_stars_repo_name: bealwang/MegEngine

max_issues_repo_name: bealwang/MegEngine

max_forks_repo_name: bealwang/MegEngine