Method forward

axlearn/common/summary_test.py:246–257

(self)

            """A test `Module`."""

            def forward(self):
                # This logs a 7 row table with two columns where each cell contains a 16 x 16 color
                # image.
                # Shape: num examples x table columns x image height x image width x channels.
                images = jax.numpy.ones((7, 2, 16, 16, 3))

                def create_table(images: np.ndarray):
                    return wandb.Table(
                        ["output", "target"], [[wandb.Image(img) for img in row] for row in images]
                    )

                self.add_summary("my_table", CallbackSummary(create_table, images))

        module = _TestModule.default_config().set(name="tmp").instantiate(parent=None)