Update b3d to GenJAX 0.6.1 (GEN-561) (#158)

sritchie · web-flow · commit 007915bfd5b5 · 2024-09-10T21:03:55.000-04:00
This PR: - upgrades GenJAX to 0.6.1 - massages a few `ChoiceMap` construction calls to use nicer, newer methods - replaces broken `.const` calls on `Const` objects with `unwrap()`. `val` would work too, but this is the preferred way. See the release notes here: https://github.com/probcomp/genjax/releases/tag/v0.6.0 and here because I fumbled the release and forgot a commit: https://github.com/probcomp/genjax/releases/tag/v0.6.1
diff --git a/notebooks/aug1demos/slam_color_room.ipynb b/notebooks/aug1demos/slam_color_room.ipynb
@@ -1621,7 +1621,7 @@
     "# convert_rgbd_to_color_space = lambda x: b3d.colors.rgbd_to_labd(x)\n",
     "# convert_color_space_to_rgbd = lambda x: b3d.colors.labd_to_rgbd(x)\n",
     "def intermediate_likelihood_func(observed_rgbd, latent_rgbd, likelihood_args):\n",
-    "    k = likelihood_args[\"k\"].const\n",
+    "    k = likelihood_args[\"k\"].unwrap()\n",
     "    fx = likelihood_args[\"fx\"]\n",
     "    fy = likelihood_args[\"fy\"]\n",
     "\n",
diff --git a/notebooks/bayes3d_paper/interactive.ipynb b/notebooks/bayes3d_paper/interactive.ipynb
@@ -169,7 +169,7 @@
                 "# def sample_likelihood_func()\n",
                 "\n",
                 "def intermediate_likelihood_func(observed_rgbd, latent_rgbd, likelihood_args):\n",
-                "    k = likelihood_args[\"k\"].const\n",
+                "    k = likelihood_args[\"k\"].unwrap()\n",
                 "    fx = likelihood_args[\"fx\"]\n",
                 "    fy = likelihood_args[\"fy\"]\n",
                 "    \n",
@@ -628,16 +628,16 @@
                 "    rendered_rgbd = renderer.render_rgbd_from_mesh(meshes[IDX].transform(pose))\n",
                 "    rendered_color_space_d = convert_rgbd_to_color_space(rendered_rgbd)\n",
                 "\n",
-                "    k = likelihood_args[\"k\"].const\n",
+                "    k = likelihood_args[\"k\"].unwrap()\n",
                 "    image_height, image_width = rendered_color_space_d.shape[0], rendered_color_space_d.shape[1]\n",
                 "    image_height = Pytree.const(image_height)\n",
                 "    image_width = Pytree.const(image_width)\n",
                 "\n",
                 "    row_coordinates = genjax.categorical.vmap(in_axes=(0,))(\n",
-                "            jnp.ones((k, image_height.const))\n",
+                "            jnp.ones((k, image_height.unwrap()))\n",
                 "    )  @ \"row_coordinates\"\n",
                 "    column_coordinates = genjax.categorical.vmap(in_axes=(0,))(\n",
-                "            jnp.ones((k, image_width.const))\n",
+                "            jnp.ones((k, image_width.unwrap()))\n",
                 "    ) @ \"column_coordinates\"\n",
                 "\n",
                 "\n",
diff --git a/notebooks/bayes3d_paper/run_ycbv_evaluation.py b/notebooks/bayes3d_paper/run_ycbv_evaluation.py
@@ -108,11 +108,11 @@ def _gvmf_and_select_best_move(
                         Pose.sample_gaussian_vmf_pose, in_axes=(0, None, None, None)
                     )(
                         jax.random.split(key, number),
-                        trace.get_choices()[address.const],
+                        trace.get_choices()[address.unwrap()],
                         variance,
                         concentration,
                     ),
-                    trace.get_choices()[address.const][None, ...],
+                    trace.get_choices()[address.unwrap()][None, ...],
                 ]
             )
             scores = jax.vmap(update_pose_and_color, in_axes=(None, None, 0))(
diff --git a/notebooks/bayes3d_paper/ycbv.ipynb b/notebooks/bayes3d_paper/ycbv.ipynb
@@ -346,7 +346,7 @@
     "@genjax.gen\n",
     "def dense_multiobject_model(num_objects, meshes, likelihood_args):\n",
     "    all_poses = []\n",
-    "    for i in range(num_objects.const):\n",
+    "    for i in range(num_objects.unwrap()):\n",
     "        object_pose = (\n",
     "            uniform_pose(jnp.ones(3) * -100.0, jnp.ones(3) * 100.0) @ f\"object_pose_{i}\"\n",
     "        )\n",
diff --git a/notebooks/integration.ipynb b/notebooks/integration.ipynb
@@ -534,7 +534,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "scantr = ps.particle_system_state_step.scan(n=(num_timesteps.const - 1)).simulate(\n",
+    "scantr = ps.particle_system_state_step.scan(n=(num_timesteps.unwrap() - 1)).simulate(\n",
     "    key, (state0, None)\n",
     ")"
    ]
diff --git a/pixi.lock b/pixi.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -109,7 +109,7 @@ cwd = "scripts"
 [tool.pixi.feature.core.pypi-dependencies]
 carvekit = "==4.1.2"
 datasync = "==0.0.2"
-genjax = "==0.5.1"
+genjax = "==0.6.1"
 pykitti = "==0.3.1"
 pyliblzfse = { git = "https://github.com/ydkhatri/pyliblzfse.git" }
 pyransac3d = ">=0.6.0,<0.7"
diff --git a/src/b3d/chisight/dense/dense_model.py b/src/b3d/chisight/dense/dense_model.py
@@ -16,8 +16,8 @@ def make_dense_multiobject_model(renderer, likelihood_func, sample_func=None):
         def f(key, likelihood_args):
             return jnp.zeros(
                 (
-                    likelihood_args["image_height"].const,
-                    likelihood_args["image_width"].const,
+                    likelihood_args["image_height"].unwrap(),
+                    likelihood_args["image_width"].unwrap(),
                     4,
                 )
             )
@@ -44,7 +44,7 @@ def dense_multiobject_model(args_dict):
         likelihood_args["blur"] = blur
 
         all_poses = []
-        for i in range(num_objects.const):
+        for i in range(num_objects.unwrap()):
             object_pose = (
                 uniform_pose(jnp.ones(3) * -100.0, jnp.ones(3) * 100.0)
                 @ f"object_pose_{i}"
diff --git a/src/b3d/chisight/dense/likelihoods/blur_likelihood_gaussian.py b/src/b3d/chisight/dense/likelihoods/blur_likelihood_gaussian.py
@@ -114,7 +114,7 @@ def likelihood_per_pixel(latent_rgbd: jnp.ndarray, blur):
 
 @jax.jit
 def blur_intermediate_likelihood_func(observed_rgbd, latent_rgbd, likelihood_args):
-    # k = likelihood_args["k"].const
+    # k = likelihood_args["k"].unwrap()
     color_variance = likelihood_args["color_variance_0"]
     depth_variance = likelihood_args["depth_variance_0"]
     outlier_probability = likelihood_args["outlier_probability_0"]
diff --git a/src/b3d/chisight/dynamic_object_model/likelihoods/aggreate_mean_image_kernel.py b/src/b3d/chisight/dynamic_object_model/likelihoods/aggreate_mean_image_kernel.py
@@ -24,7 +24,7 @@ def sample_func(key, args):
     ).astype(jnp.int32)
 
     latent_image_sum = jnp.zeros(
-        (args["image_height"].const, args["image_width"].const, 4)
+        (args["image_height"].unwrap(), args["image_width"].unwrap(), 4)
     )
     latent_image_sum = latent_image_sum.at[pixels[..., 0], pixels[..., 1], :3].add(
         args["colors"] * (1 - args["color_outlier_probability"])[:, None]
@@ -34,14 +34,14 @@ def sample_func(key, args):
     )
 
     projected_points_count = jnp.zeros(
-        (args["image_height"].const, args["image_width"].const)
+        (args["image_height"].unwrap(), args["image_width"].unwrap())
     )
     projected_points_count = projected_points_count.at[
         pixels[..., 0], pixels[..., 1]
     ].add(1)
 
     non_registration_probability = jnp.ones(
-        (args["image_height"].const, args["image_width"].const)
+        (args["image_height"].unwrap(), args["image_width"].unwrap())
     )
     non_registration_probability = non_registration_probability.at[
         pixels[..., 0], pixels[..., 1]
@@ -81,7 +81,7 @@ def likelihood_func(observed_rgbd, args):
     ).astype(jnp.int32)
 
     latent_image_sum = jnp.zeros(
-        (args["image_height"].const, args["image_width"].const, 4)
+        (args["image_height"].unwrap(), args["image_width"].unwrap(), 4)
     )
     latent_image_sum = latent_image_sum.at[pixels[..., 0], pixels[..., 1], :3].add(
         args["colors"] * (1 - args["color_outlier_probability"])[:, None]
@@ -91,14 +91,14 @@ def likelihood_func(observed_rgbd, args):
     )
 
     projected_points_count = jnp.zeros(
-        (args["image_height"].const, args["image_width"].const)
+        (args["image_height"].unwrap(), args["image_width"].unwrap())
     )
     projected_points_count = projected_points_count.at[
         pixels[..., 0], pixels[..., 1]
     ].add(1)
 
     non_registration_probability = jnp.ones(
-        (args["image_height"].const, args["image_width"].const)
+        (args["image_height"].unwrap(), args["image_width"].unwrap())
     )
     non_registration_probability = non_registration_probability.at[
         pixels[..., 0], pixels[..., 1]
diff --git a/src/b3d/chisight/dynamic_object_model/likelihoods/project_no_occlusions_kernel.py b/src/b3d/chisight/dynamic_object_model/likelihoods/project_no_occlusions_kernel.py
@@ -65,8 +65,8 @@ def likelihood_func(observed_rgbd, args):
 def sample_func(key, likelihood_args):
     return jnp.zeros(
         (
-            likelihood_args["image_height"].const,
-            likelihood_args["image_width"].const,
+            likelihood_args["image_height"].unwrap(),
+            likelihood_args["image_width"].unwrap(),
             4,
         )
     )
diff --git a/src/b3d/chisight/particle_system.py b/src/b3d/chisight/particle_system.py
@@ -26,22 +26,22 @@ def initial_particle_system_state(
 ):
     relative_particle_poses = (
         dummy_mapped_uniform_pose(
-            jnp.arange(num_particles.const), *relative_particle_poses_prior_params
+            jnp.arange(num_particles.unwrap()), *relative_particle_poses_prior_params
         )
         @ "particle_poses"
     )
 
     object_assignments = (
         b3d.modeling_utils.categorical.vmap(in_axes=(0,))(
-            jnp.zeros((num_particles.const, num_clusters.const))
+            jnp.zeros((num_particles.unwrap(), num_clusters.unwrap()))
         )
         @ "object_assignments"
     )
 
     # Cluster pose in world coordinates
     initial_object_poses = (
         dummy_mapped_uniform_pose(
-            jnp.arange(num_clusters.const), *initial_object_poses_prior_params
+            jnp.arange(num_clusters.unwrap()), *initial_object_poses_prior_params
         )
         @ "object_poses"
     )
@@ -59,7 +59,7 @@ def initial_particle_system_state(
     # Initial visibility mask
     initial_vis_mask = (
         b3d.modeling_utils.bernoulli.vmap(in_axes=(0,))(
-            jnp.repeat(jax.scipy.special.logit(0.5), num_particles.const)
+            jnp.repeat(jax.scipy.special.logit(0.5), num_particles.unwrap())
         )
         @ "initial_visibility"
     )
@@ -101,7 +101,7 @@ def particle_system_state_step(carried_state, _):
     # Visibility mask
     vis_mask = (
         b3d.modeling_utils.bernoulli.vmap(in_axes=(0,))(
-            jnp.repeat(jax.scipy.special.logit(0.5), num_particles.const)
+            jnp.repeat(jax.scipy.special.logit(0.5), num_particles.unwrap())
         )
         @ "visibility"
     )
@@ -146,7 +146,7 @@ def latent_particle_model(
     )
 
     final_state, scan_retvals = (
-        particle_system_state_step.scan(n=(num_timesteps.const - 1))(state0, None)
+        particle_system_state_step.scan(n=(num_timesteps.unwrap() - 1))(state0, None)
         @ "states1+"
     )
 
@@ -164,7 +164,7 @@ def sparse_observation_model(
 ):
     # TODO: add visibility
     uv = b3d.camera.screen_from_world(
-        particle_absolute_poses.pos, camera_pose, instrinsics.const
+        particle_absolute_poses.pos, camera_pose, instrinsics.unwrap()
     )
     uv_ = (
         b3d.modeling_utils.normal(uv, jnp.tile(sigma, uv.shape)) @ "sensor_coordinates"
@@ -248,7 +248,7 @@ def visualize_particle_system(
     camera_pose = particle_dynamics_summary["camera_pose"]
     object_assignments = static_state[0]
 
-    cluster_colors = jnp.array(b3d.distinct_colors(num_clusters.const))
+    cluster_colors = jnp.array(b3d.distinct_colors(num_clusters.unwrap()))
 
     rr.log(
         f"{viz_prefix}/3D",
@@ -266,7 +266,7 @@ def visualize_particle_system(
         timeless=True,
     )
 
-    for t in range(num_timesteps.const):
+    for t in range(num_timesteps.unwrap()):
         rr.set_time_sequence("time", t)
 
         cam_pose = camera_pose[t]
@@ -286,7 +286,7 @@ def visualize_particle_system(
             ),
         )
 
-        for i in range(num_clusters.const):
+        for i in range(num_clusters.unwrap()):
             b3d.rr_log_pose(f"{viz_prefix}/3D/cluster/{i}", object_poses[t][i])
 
 
@@ -302,7 +302,7 @@ def particle_2d_pixel_coordinates_to_image(pixel_coords, image_height, image_wid
 def visualize_sparse_observation(sparse_model_args, observations):
     import rerun as rr
 
-    intrinsics = sparse_model_args[0].const
+    intrinsics = sparse_model_args[0].unwrap()
 
     for t in range(observations.shape[0]):
         rr.set_time_sequence("time", t)
@@ -338,7 +338,7 @@ def visualize_dense_gps(
             timeless=True,
         )
 
-    for t in range(num_timesteps.const):
+    for t in range(num_timesteps.unwrap()):
         rr.set_time_sequence("time", t)
         poses = particle_dynamics_summary["absolute_particle_poses"][t]
         for i in range(len(meshes)):
diff --git a/src/b3d/chisight/patch_tracking.py b/src/b3d/chisight/patch_tracking.py
@@ -90,7 +90,7 @@ def get_adam_optimization_patch_tracker(model, patches, pose_WC=Pose.identity())
     """
 
     def allidx_chm(x):
-        return genjax.ChoiceMap.idx(jnp.arange(x.shape[0], dtype=int), x)
+        return C[jnp.arange(x.shape[0], dtype=int)].set(x)
 
     @jax.jit
     def importance_from_pos_quat(positions, quaternions, observed_rgbd):
@@ -118,14 +118,14 @@ def importance_from_pos_quat(positions, quaternions, observed_rgbd):
         )
 
         particle_poses = jax.tree.map(
-            lambda arr: jnp.tile(arr, (num_particles.const, 1)), Pose.identity()
+            lambda arr: jnp.tile(arr, (num_particles.unwrap(), 1)), Pose.identity()
         )
-        object_assignments = jnp.arange(num_particles.const, dtype=int)
+        object_assignments = jnp.arange(num_particles.unwrap(), dtype=int)
         object_poses = jax.vmap(
             lambda pos, quat: Pose.from_vec(jnp.concatenate([pos, quat])),
             in_axes=(0, 0),
         )(positions, quaternions)
-        vis_mask = jnp.ones((num_particles.const,), dtype=int)
+        vis_mask = jnp.ones((num_particles.unwrap(),), dtype=int)
 
         constraints = C.d(
             {
diff --git a/src/b3d/utils.py b/src/b3d/utils.py
@@ -511,7 +511,7 @@ def multivmap(f, args=None):
 def update_choices(trace, addresses, *values):
     return trace.update(
         jax.random.PRNGKey(0),
-        genjax.ChoiceMap.d({addr: c for (addr, c) in zip(addresses.const, values)}),
+        genjax.ChoiceMap.from_mapping(zip(addresses.unwrap(), values)),
     )[0]
 
 
@@ -549,13 +549,13 @@ def update_choices_get_score(trace, addr_const, *values):
 
 @jax.jit
 def grid_trace(trace, addresses_const, values):
-    if len(addresses_const.const) == 1:
+    if len(addresses_const.unwrap()) == 1:
         return grid1(trace, addresses_const, *values)
-    elif len(addresses_const.const) == 2:
+    elif len(addresses_const.unwrap()) == 2:
         return grid2(trace, addresses_const, *values)
-    elif len(addresses_const.const) == 3:
+    elif len(addresses_const.unwrap()) == 3:
         return grid3(trace, addresses_const, *values)
-    elif len(addresses_const.const) == 4:
+    elif len(addresses_const.unwrap()) == 4:
         return grid4(trace, addresses_const, *values)
     else:
         raise ValueError("Too many addresses")
diff --git a/tests/dense_model_unit_tests/triangle_depth_posterior/solver/importance.py b/tests/dense_model_unit_tests/triangle_depth_posterior/solver/importance.py
@@ -67,8 +67,8 @@ def importance_sample_with_depth_in_partition(
         C.d(
             {
                 "triangle_vertices": new_triangle_W,
-                "observed_rgbs": genjax.ChoiceMap.idx(
-                    jnp.arange(T), C.n().at["observed_rgb"].set(task_input["video"])
+                "observed_rgbs": C[jnp.arange(T), "observed_rgb"].set(
+                    task_input["video"]
                 ),
             }
         ),
diff --git a/tests/test_chisight_dense_gps.py b/tests/test_chisight_dense_gps.py
@@ -53,8 +53,8 @@ def cube_mesh_with_size_and_color(size, color):
         return mesh
 
     meshes = jax.vmap(cube_mesh_with_size_and_color)(
-        jnp.ones((num_particles.const, 3)) * jnp.array([[0.1, 0.1, 0.01]]),
-        jax.random.uniform(key, (num_particles.const, 3)),
+        jnp.ones((num_particles.unwrap(), 3)) * jnp.array([[0.1, 0.1, 0.01]]),
+        jax.random.uniform(key, (num_particles.unwrap(), 3)),
     )
 
     likelihood = make_krays_image_observation_model(renderer)

Original file line number	Diff line number	Diff line change
`@@ -108,11 +108,11 @@ def _gvmf_and_select_best_move(`
`108`	`108`	`Pose.sample_gaussian_vmf_pose, in_axes=(0, None, None, None)`
`109`	`109`	`)(`
`110`	`110`	`jax.random.split(key, number),`
`111`		`- trace.get_choices()[address.const],`
	`111`	`+ trace.get_choices()[address.unwrap()],`
`112`	`112`	`variance,`
`113`	`113`	`concentration,`
`114`	`114`	`),`
`115`		`- trace.get_choices()[address.const][None, ...],`
	`115`	`+ trace.get_choices()[address.unwrap()][None, ...],`
`116`	`116`	`]`
`117`	`117`	`)`
`118`	`118`	`scores = jax.vmap(update_pose_and_color, in_axes=(None, None, 0))(`
Original file line number	Diff line number	Diff line change
`@@ -534,7 +534,7 @@`
`534`	`534`	`"metadata": {},`
`535`	`535`	`"outputs": [],`
`536`	`536`	`"source": [`
`537`		`- "scantr = ps.particle_system_state_step.scan(n=(num_timesteps.const - 1)).simulate(\n",`
	`537`	`+ "scantr = ps.particle_system_state_step.scan(n=(num_timesteps.unwrap() - 1)).simulate(\n",`
`538`	`538`	`" key, (state0, None)\n",`
`539`	`539`	`")"`
`540`	`540`	`]`
Original file line number	Diff line number	Diff line change
`@@ -16,8 +16,8 @@ def make_dense_multiobject_model(renderer, likelihood_func, sample_func=None):`
`16`	`16`	`def f(key, likelihood_args):`
`17`	`17`	`return jnp.zeros(`
`18`	`18`	`(`
`19`		`- likelihood_args["image_height"].const,`
`20`		`- likelihood_args["image_width"].const,`
	`19`	`+ likelihood_args["image_height"].unwrap(),`
	`20`	`+ likelihood_args["image_width"].unwrap(),`
`21`	`21`	`4,`
`22`	`22`	`)`
`23`	`23`	`)`
`@@ -44,7 +44,7 @@ def dense_multiobject_model(args_dict):`
`44`	`44`	`likelihood_args["blur"] = blur`
`45`	`45`
`46`	`46`	`all_poses = []`
`47`		`- for i in range(num_objects.const):`
	`47`	`+ for i in range(num_objects.unwrap()):`
`48`	`48`	`object_pose = (`
`49`	`49`	`uniform_pose(jnp.ones(3) * -100.0, jnp.ones(3) * 100.0)`
`50`	`50`	`@ f"object_pose_{i}"`
Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,7 @@ def sample_func(key, args):`
`24`	`24`	`).astype(jnp.int32)`
`25`	`25`
`26`	`26`	`latent_image_sum = jnp.zeros(`
`27`		`- (args["image_height"].const, args["image_width"].const, 4)`
	`27`	`+ (args["image_height"].unwrap(), args["image_width"].unwrap(), 4)`
`28`	`28`	`)`
`29`	`29`	`latent_image_sum = latent_image_sum.at[pixels[..., 0], pixels[..., 1], :3].add(`
`30`	`30`	`args["colors"] * (1 - args["color_outlier_probability"])[:, None]`
`@@ -34,14 +34,14 @@ def sample_func(key, args):`
`34`	`34`	`)`
`35`	`35`
`36`	`36`	`projected_points_count = jnp.zeros(`
`37`		`- (args["image_height"].const, args["image_width"].const)`
	`37`	`+ (args["image_height"].unwrap(), args["image_width"].unwrap())`
`38`	`38`	`)`
`39`	`39`	`projected_points_count = projected_points_count.at[`
`40`	`40`	`pixels[..., 0], pixels[..., 1]`
`41`	`41`	`].add(1)`
`42`	`42`
`43`	`43`	`non_registration_probability = jnp.ones(`
`44`		`- (args["image_height"].const, args["image_width"].const)`
	`44`	`+ (args["image_height"].unwrap(), args["image_width"].unwrap())`
`45`	`45`	`)`
`46`	`46`	`non_registration_probability = non_registration_probability.at[`
`47`	`47`	`pixels[..., 0], pixels[..., 1]`
`@@ -81,7 +81,7 @@ def likelihood_func(observed_rgbd, args):`
`81`	`81`	`).astype(jnp.int32)`
`82`	`82`
`83`	`83`	`latent_image_sum = jnp.zeros(`
`84`		`- (args["image_height"].const, args["image_width"].const, 4)`
	`84`	`+ (args["image_height"].unwrap(), args["image_width"].unwrap(), 4)`
`85`	`85`	`)`
`86`	`86`	`latent_image_sum = latent_image_sum.at[pixels[..., 0], pixels[..., 1], :3].add(`
`87`	`87`	`args["colors"] * (1 - args["color_outlier_probability"])[:, None]`
`@@ -91,14 +91,14 @@ def likelihood_func(observed_rgbd, args):`
`91`	`91`	`)`
`92`	`92`
`93`	`93`	`projected_points_count = jnp.zeros(`
`94`		`- (args["image_height"].const, args["image_width"].const)`
	`94`	`+ (args["image_height"].unwrap(), args["image_width"].unwrap())`
`95`	`95`	`)`
`96`	`96`	`projected_points_count = projected_points_count.at[`
`97`	`97`	`pixels[..., 0], pixels[..., 1]`
`98`	`98`	`].add(1)`
`99`	`99`
`100`	`100`	`non_registration_probability = jnp.ones(`
`101`		`- (args["image_height"].const, args["image_width"].const)`
	`101`	`+ (args["image_height"].unwrap(), args["image_width"].unwrap())`
`102`	`102`	`)`
`103`	`103`	`non_registration_probability = non_registration_probability.at[`
`104`	`104`	`pixels[..., 0], pixels[..., 1]`