File size: 345 Bytes
3d7c303
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
import numpy as np
from pysr import sympy2torch
import torch
import sympy

print("Test Torch 1 - test export")
x, y, z = sympy.symbols('x y z')
cosx = 1.0 * sympy.cos(x) + y
X = torch.randn((1000, 2))
true = 1.0 * torch.cos(X[:, 0]) + X[:, 1]
torch_module = sympy2torch(cosx, [x, y, z])
assert jnp.all(jnp.isclose(torch_module(X), true)).item()