import torch
x = torch.tensor([1., 2., 3., 4., 5., 6.])

import tensorflow as tf
x = tf.constant([1., 2., 3., 4., 5., 6.])

import jax.numpy as np
x = np.array([1., 2., 3., 4., 5., 6.])

import numpy as np
x = np.array([1., 2., 3., 4., 5., 6.])

# --------------------------------------------------------------------------

# No matter which framwork you use, you can use the same code
import eagerpy as ep

# Just wrap a native tensor using EagerPy
x = ep.astensor(x)

# All of EagerPy's functionality is available as methods ...
x = x.reshape((2, 3))
norms = x.flatten(start=1).square().sum(axis=-1).sqrt()
norms = x.flatten(start=1)

# ... and functions
_, grad = ep.value_and_grad(loss_fn, x)
ep.clip(x + eps * grad, 0, 1)

# You can even write functions that work transparently with
# Pytorch tensors, TensorFlow tensors, JAX arrays, NumPy arrays
# and EagerPy tensors
def squared_a_plus_b_times_c(a, b, c):
   (a, b, c), restore_type = ep.astensors_(a, b, c)
   # here, a, b, c are EagerPyTensors
   result = (a + b * c).square()
   return restore_type(result)

# You can call this function using any kind of tensors and the result
# will have the same type.