.. raw:: html

.. image:: https://badge.fury.io/py/eagerpy.svg :target: https://badge.fury.io/py/eagerpy

.. image:: https://codecov.io/gh/jonasrauber/eagerpy/branch/master/graph/badge.svg :target: https://codecov.io/gh/jonasrauber/eagerpy

.. image:: https://img.shields.io/badge/code%20style-black-000000.svg :target: https://github.com/ambv/black

================================================================================== EagerPy：编写可原生兼容 PyTorch、TensorFlow、JAX 和 NumPy 的代码

EagerPy <https://eagerpy.jonasrauber.de>_ 是一个 Python 框架，它使您能够编写自动原生兼容 PyTorch <https://pytorch.org>、TensorFlow <https://www.tensorflow.org>、JAX <https://github.com/google/jax>_ 和 NumPy <https://numpy.org>_ 的代码。即使您只使用其中一个框架，EagerPy 也非常出色，因为它提供了一个简洁一致的 API，完全支持链式调用，具有丰富的类型注解，并能让您写出优雅的代码。

🔥 设计目标

• 原生性能：EagerPy 的操作会直接转换为相应的原生操作。
• 完全可链式调用：所有功能既可以通过张量对象的方法调用，也可以通过 EagerPy 函数调用。
• 类型检查：借助 EagerPy 丰富的类型注解，在运行代码之前就能捕获错误。

📖 文档

在 文档 <https://eagerpy.jonasrauber.de>_ 中了解更多关于 EagerPy 的信息。

🚀 快速入门

.. code-block:: bash

pip install eagerpy

EagerPy 需要 Python 3.6 或更高版本。除此之外，所有必要的依赖项都会自动安装。若要与 PyTorch、TensorFlow、JAX 或 NumPy 一起使用，需要单独安装相应的框架。这些框架并未被声明为依赖项，因为并非所有人都希望同时使用并安装它们，而且其中一些包针对不同的架构和 CUDA <https://developer.nvidia.com/cuda-zone>_ 版本有不同的构建版本。

🎉 示例

.. code-block:: python

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.])

无论使用哪个框架，都可以使用相同的代码

import eagerpy as ep

只需使用 EagerPy 包装原生张量

x = ep.astensor(x)

EagerPy 的所有功能都可以作为方法调用

x = x.reshape((2, 3)) x.flatten(start=1).square().sum(axis=-1).sqrt()

或者：x.flatten(1).norms.l2()

也可以作为函数调用（是的，梯度也受支持！）

loss, grad = ep.value_and_grad(loss_fn, x) ep.clip(x + eps * grad, 0, 1)

您甚至可以编写对 Pytorch 张量、TensorFlow 张量、JAX 数组、NumPy 数组都透明兼容的函数

def my_universal_function(a, b, c): # 将所有输入转换为 EagerPy 张量 a, b, c = ep.astensors(a, b, c)

   # 执行一些计算
   result = (a + b * c).square()

   # 并返回原生张量
   return result.raw

🗺 使用场景

Foolbox Native <https://github.com/bethgelab/foolbox>_，即广受欢迎的对抗攻击库 Foolbox 的最新版本，已使用 EagerPy 而不是 NumPy 从头重写，从而在基于 PyTorch、TensorFlow 和 JAX 构建的模型上实现了原生性能，且仅需一套代码库。

EagerPy 还被其他框架用于减少代码重复（例如 GUDHI <https://github.com/GUDHI/gudhi-devel>），或用于 比较不同框架的性能 <https://github.com/jonasrauber/uniformly-sampling-nd-ball>。

📄 引用

如果您使用 EagerPy，请使用以下 BibTex 条目引用我们的论文 arXiv:2008.04175 <https://arxiv.org/abs/2008.04175>_：

.. code-block::

@article{rauber2020eagerpy, title={{EagerPy}: Writing Code That Works Natively with {PyTorch}, {TensorFlow}, {JAX}, and {NumPy}}, author={Rauber, Jonas and Bethge, Matthias and Brendel, Wieland}, journal={arXiv preprint arXiv:2008.04175}, year={2020}, url={https://eagerpy.jonasrauber.de}, }

🐍 兼容性

我们目前测试的版本如下：

• PyTorch 1.4.0
• TensorFlow 2.1.0
• JAX 0.1.57
• NumPy 1.18.1

EagerPy 快速上手指南

EagerPy 是一个 Python 框架，让你编写的代码能够原生地同时适用于 PyTorch、TensorFlow、JAX 和 NumPy。它提供完全可链式调用（Chainable）的 API 和丰富的类型注解，帮助你编写简洁、高效且跨框架兼容的代码。

环境准备

• 系统要求：Python 3.6 或更高版本。
• 前置依赖：
  • EagerPy 会自动安装其核心依赖。
  • 注意：你需要单独安装目标深度学习框架（PyTorch、TensorFlow、JAX 或 NumPy 中的至少一个）。EagerPy 未将它们列为强制依赖，以便用户根据硬件架构（如 CUDA 版本）灵活选择安装特定版本。

安装步骤

使用 pip 安装 EagerPy：

pip install eagerpy

提示：国内开发者可使用清华或阿里镜像源加速安装：

pip install eagerpy -i https://pypi.tuna.tsinghua.edu.cn/simple

安装完成后，请确保已单独安装你需要的框架，例如：

pip install torch tensorflow jax numpy

基本使用

EagerPy 的核心用法是将不同框架的原生张量（Tensor）包装为 EagerPy 张量，之后即可使用统一的链式方法或函数进行操作。

1. 包装原生张量

无论底层使用的是哪个框架，都可以通过 ep.astensor() 进行统一包装：

import torch
import tensorflow as tf
import jax.numpy as np
import numpy as np
import eagerpy as ep

# 假设这是来自任意框架的原生张量
x_native = torch.tensor([1., 2., 3., 4., 5., 6.]) 
# 或者: x_native = tf.constant([...])
# 或者: x_native = np.array([...])

# 包装为 EagerPy 张量
x = ep.astensor(x_native)

2. 链式调用与方法

所有功能均可作为张量对象的方法直接调用，支持流畅的链式操作：

# 重塑、展平、平方、求和、开方等操作
result = x.reshape((2, 3)).flatten(start=1).square().sum(axis=-1).sqrt()

# 简化写法：计算 L2 范数
norm_result = x.flatten(1).norms.l2()

3. 函数式调用与梯度计算

EagerPy 也提供函数式接口，并原生支持自动微分（梯度计算）：

def loss_fn(x):
    return x.square().sum()

# 同时获取损失值和梯度
loss, grad = ep.value_and_grad(loss_fn, x)

# 对数据进行裁剪操作
clipped_x = ep.clip(x + 0.1 * grad, 0, 1)

4. 编写通用函数

你可以编写一次代码，使其透明地适配所有支持的框架：

def my_universal_function(a, b, c):
    # 将所有输入转换为 EagerPy 张量
    a, b, c = ep.astensors(a, b, c)
    
    # 执行计算（链式或常规运算均可）
    result = (a + b * c).square()
    
    # 返回原生张量（通过 .raw 属性）
    return result.raw