Remove references to deprecated submodule jax.abstract_arrays

Use jax.core instead (see jax-ml/jax#16271) PiperOrigin-RevId: 538767153
google-research · Jun 8, 2023 · 9883169 · 9883169
1 parent 7f1e018
commit 9883169
Show file tree

Hide file tree

Showing 48 changed files with 53 additions and 54 deletions.
diff --git a/diffren/__init__.py b/diffren/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/common/__init__.py b/diffren/common/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/common/compare_images.py b/diffren/common/compare_images.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/common/kernels/rasterize_test_utils.cc b/diffren/common/kernels/rasterize_test_utils.cc
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/common/kernels/rasterize_test_utils.h b/diffren/common/kernels/rasterize_test_utils.h
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/common/kernels/rasterize_triangles_impl.cc b/diffren/common/kernels/rasterize_triangles_impl.cc
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/common/kernels/rasterize_triangles_impl.h b/diffren/common/kernels/rasterize_triangles_impl.h
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/common/kernels/rasterize_triangles_impl_cuda.cu.cc b/diffren/common/kernels/rasterize_triangles_impl_cuda.cu.cc
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/common/kernels/rasterize_triangles_impl_cuda.cu.h b/diffren/common/kernels/rasterize_triangles_impl_cuda.cu.h
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/common/kernels/rasterize_triangles_impl_cuda_test.cc b/diffren/common/kernels/rasterize_triangles_impl_cuda_test.cc
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/common/kernels/rasterize_triangles_impl_test.cc b/diffren/common/kernels/rasterize_triangles_impl_test.cc
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/common/kernels/rasterize_triangles_types.h b/diffren/common/kernels/rasterize_triangles_types.h
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/common/kernels/rasterize_utils.h b/diffren/common/kernels/rasterize_utils.h
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/common/obj_loader.py b/diffren/common/obj_loader.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/common/obj_loader_test.py b/diffren/common/obj_loader_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/common/test_utils.py b/diffren/common/test_utils.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/__init__.py b/diffren/jax/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/camera.py b/diffren/jax/camera.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/camera_test.py b/diffren/jax/camera_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/composite.py b/diffren/jax/composite.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/constants.py b/diffren/jax/constants.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/__init__.py b/diffren/jax/internal/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/framebuffer.py b/diffren/jax/internal/framebuffer.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/framebuffer_test.py b/diffren/jax/internal/framebuffer_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/interpolate.py b/diffren/jax/internal/interpolate.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/interpolate_test.py b/diffren/jax/internal/interpolate_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/kernels/descriptors.proto b/diffren/jax/internal/kernels/descriptors.proto
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/kernels/rasterize_triangles_cpu.cc b/diffren/jax/internal/kernels/rasterize_triangles_cpu.cc
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/kernels/rasterize_triangles_gpu.cc b/diffren/jax/internal/kernels/rasterize_triangles_gpu.cc
@@ -1,4 +1,4 @@
-/* Copyright 2022 The diffren Authors.
+/* Copyright 2023 The diffren Authors.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/kernels/rasterize_triangles_xla.py b/diffren/jax/internal/kernels/rasterize_triangles_xla.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,7 +20,6 @@
 from . import rasterize_triangles_gpu
 from diffren.jax.internal.kernels import descriptors_pb2
 import jax
-from jax import abstract_arrays
 from jax.interpreters import batching
 import jax.numpy as jnp
 import numpy as np
@@ -150,11 +149,11 @@ def rasterize_triangles_abstract_eval(vertices, triangles, *, image_width,
   assert triangles.shape[1] == 3
   del face_culling_mode
 
-  return (abstract_arrays.ShapedArray((num_layers, image_height, image_width),
-                                      dtype=np.int32),
-          abstract_arrays.ShapedArray((num_layers, image_height, image_width),
-                                      dtype=np.float32),
-          abstract_arrays.ShapedArray(
+  return (jax.core.ShapedArray((num_layers, image_height, image_width),
+                               dtype=np.int32),
+          jax.core.ShapedArray((num_layers, image_height, image_width),
+                               dtype=np.float32),
+          jax.core.ShapedArray(
               (num_layers, image_height, image_width, 3), dtype=np.float32))
 
 

diff --git a/diffren/jax/internal/kernels/rasterize_triangles_xla_test.py b/diffren/jax/internal/kernels/rasterize_triangles_xla_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/rasterize.py b/diffren/jax/internal/rasterize.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/rasterize_test.py b/diffren/jax/internal/rasterize_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/splat.py b/diffren/jax/internal/splat.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/internal/splat_test.py b/diffren/jax/internal/splat_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/render.py b/diffren/jax/render.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/render_test.py b/diffren/jax/render_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/__init__.py b/diffren/jax/utils/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/depthmap.py b/diffren/jax/utils/depthmap.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/depthmap_test.py b/diffren/jax/utils/depthmap_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/image.py b/diffren/jax/utils/image.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/image_test.py b/diffren/jax/utils/image_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/mesh.py b/diffren/jax/utils/mesh.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/mesh_test.py b/diffren/jax/utils/mesh_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/shaders.py b/diffren/jax/utils/shaders.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/shaders_test.py b/diffren/jax/utils/shaders_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/transforms.py b/diffren/jax/utils/transforms.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/diffren/jax/utils/transforms_test.py b/diffren/jax/utils/transforms_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The diffren Authors.
+# Copyright 2023 The diffren Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.