From be1429c2d216100f2831fcaddc4752f867a06895 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Sun, 26 Jun 2022 15:55:38 +0300
Subject: [PATCH] update numpy coverage

---
 doc/make_numpy_coverage_table.py |   2 +
 doc/numpy_coverage.rst           | 102 +++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+)

diff --git a/doc/make_numpy_coverage_table.py b/doc/make_numpy_coverage_table.py
index e100f0b3..f30d328c 100644
--- a/doc/make_numpy_coverage_table.py
+++ b/doc/make_numpy_coverage_table.py
@@ -62,7 +62,9 @@ def initialize_contexts():
 
     return [
             arraycontext.PyOpenCLArrayContext(queue, force_device_scalars=True),
+            arraycontext.EagerJAXArrayContext(),
             arraycontext.PytatoPyOpenCLArrayContext(queue),
+            arraycontext.PytatoJAXArrayContext(),
             ]
 
 
diff --git a/doc/numpy_coverage.rst b/doc/numpy_coverage.rst
index 8d58fbe7..5a7b2918 100644
--- a/doc/numpy_coverage.rst
+++ b/doc/numpy_coverage.rst
@@ -15,22 +15,34 @@ Array creation routines
 
     * - Function
       - :class:`~arraycontext.PyOpenCLArrayContext`
+      - :class:`~arraycontext.EagerJAXArrayContext`
       - :class:`~arraycontext.PytatoPyOpenCLArrayContext`
+      - :class:`~arraycontext.PytatoJAXArrayContext`
     * - :func:`numpy.empty_like`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.ones_like`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.zeros_like`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.full_like`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.copy`
+      - :green:`Yes`
       - :green:`Yes`
       - :red:`No`
+      - :red:`No`
 
 Array manipulation routines
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -40,25 +52,39 @@ Array manipulation routines
 
     * - Function
       - :class:`~arraycontext.PyOpenCLArrayContext`
+      - :class:`~arraycontext.EagerJAXArrayContext`
       - :class:`~arraycontext.PytatoPyOpenCLArrayContext`
+      - :class:`~arraycontext.PytatoJAXArrayContext`
     * - :func:`numpy.reshape`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.ravel`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.transpose`
       - :red:`No`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.broadcast_to`
       - :red:`No`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.concatenate`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.stack`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
 
 Linear algebra
 ~~~~~~~~~~~~~~
@@ -68,10 +94,14 @@ Linear algebra
 
     * - Function
       - :class:`~arraycontext.PyOpenCLArrayContext`
+      - :class:`~arraycontext.EagerJAXArrayContext`
       - :class:`~arraycontext.PytatoPyOpenCLArrayContext`
+      - :class:`~arraycontext.PytatoJAXArrayContext`
     * - :func:`numpy.vdot`
+      - :green:`Yes`
       - :green:`Yes`
       - :red:`No`
+      - :red:`No`
 
 Logic Functions
 ~~~~~~~~~~~~~~~
@@ -81,31 +111,49 @@ Logic Functions
 
     * - Function
       - :class:`~arraycontext.PyOpenCLArrayContext`
+      - :class:`~arraycontext.EagerJAXArrayContext`
       - :class:`~arraycontext.PytatoPyOpenCLArrayContext`
+      - :class:`~arraycontext.PytatoJAXArrayContext`
     * - :func:`numpy.all`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.any`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.greater`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.greater_equal`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.less`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.less_equal`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.equal`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.not_equal`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
 
 Mathematical functions
 ~~~~~~~~~~~~~~~~~~~~~~
@@ -115,82 +163,136 @@ Mathematical functions
 
     * - Function
       - :class:`~arraycontext.PyOpenCLArrayContext`
+      - :class:`~arraycontext.EagerJAXArrayContext`
       - :class:`~arraycontext.PytatoPyOpenCLArrayContext`
+      - :class:`~arraycontext.PytatoJAXArrayContext`
     * - :data:`numpy.sin`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.cos`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.tan`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.arcsin`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.arccos`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.arctan`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.arctan2`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.sinh`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.cosh`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.tanh`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.floor`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.ceil`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.sum`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.exp`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.log`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.log10`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.real`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.imag`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.conjugate`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.maximum`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.amax`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.minimum`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :func:`numpy.amin`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.sqrt`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.absolute`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`
     * - :data:`numpy.fabs`
       - :green:`Yes`
       - :green:`Yes`
+      - :green:`Yes`
+      - :green:`Yes`