Skip to content

Make debug_print more readable #313

Closed
@ricardoV94

Description

@ricardoV94
import pytensor
import pytensor.tensor as pt

x = pt.scalar("x")
y = pt.vector("y")
z = x + pt.exp(y)

pytensor.dprint(z)
Elemwise{add,no_inplace} [id A]
 |InplaceDimShuffle{x} [id B]
 | |x [id C]
 |Elemwise{exp,no_inplace} [id D]
   |y [id E]

This contains a lot of detail that is irrelevant for most users:

  1. Elemwise{...}. Almost all operations are Elemwise because users don't build graphs with scalars. I suggest we just call Add. The "type" of addition can be guessed from the input types which are visible with print_type=True. This will be even more relevant after Implement graph.vectorize and Blockwise Op #306 where everything that is not an Elemwise will likely be a Blockwise. We can maybe add a kwarg to show the exact Op that is being used.
  2. The inplace/no_inplace is seldom important for users. This information is still visible via either print_destroy_map or print_view_map. Maybe we can have a print_memory_map that is equivalent to setting both of those to True.
  3. DimShuffles are incredibly common, but very obscure because they are a term only used in PyTensor. I suggest we specialize the name for the following 4 cases:
    1. ExpandDims{axis/axes=dims}
    2. DropDims{axis/axes=dims} (or Squeeze?)
    3. Transpose{axis/axes=dims}
    4. DimShuffle{order} (if it's a mix of the ones above)

The default dprint for the above graph would then look something like:

Add [id A]
 |ExpandDims{axis=0} [id B]
 | |x [id C]
 |Exp [id D]
   |y [id E]

And dprint(print_type=True, print_memory_map=True)

Add [id A] <TensorType(float64, (?,))>
 |ExpandDims{axis=0} [id B] <TensorType(float64, (1,))> v={0: [0]}
 | |x [id C] <TensorType(float64, ())>
 |Exp [id D] <TensorType(float64, (?,))>
   |y [id E] <TensorType(float64, (?,))>

Which I think is way more readable! We can go one step further and rename <TensorType(float64, (?,))> to <Tensor(float64, shape=(?,))> or even <Vector(float64, shape=(?,))>

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions